dotnet/molilian.core/Core/taoke/AiOpenRiskControlCore.cs

using Dapper;
using dodohold.core;

namespace molilian.core
{
    public static class AiOpenRiskControlCore
    {
        public const int EndpointId = 8;
        public const string EndpointName = "aiopen";
        public const string DailyLimitReasonCode = "AI_OPEN_DAILY_LIMIT";

        public static bool IsApplicableAccount(TkPoolDTO? account)
        {
            if (account == null || string.IsNullOrWhiteSpace(account.parseEndpoint))
            {
                return false;
            }

            int[] endpointIds = account.parseEndpoint
                .Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries)
                .Select(value => int.TryParse(value, out int endpointId) ? endpointId : 0)
                .Where(endpointId => endpointId > 0)
                .Distinct()
                .ToArray();

            return endpointIds.Length == 1 && endpointIds[0] == EndpointId;
        }

        public static async Task RecordDailyLimitAndSuspendAsync(
            TkPoolDTO account,
            string reason,
            int bizErrorCode,
            int resultCode)
        {
            if (!IsApplicableAccount(account))
            {
                return;
            }

            DateTime triggerTime = DateTime.Now;
            DateTime releaseTime = triggerTime.Date.AddDays(1);
            TimeSpan suspendDuration = releaseTime - triggerTime;

            // Stop assigning this account locally before performing report aggregation.
            // The first DB writer below broadcasts the same TTL to all other nodes.
            await TkEndpointManager.SuspendForDurationAsync(
                account.id,
                EndpointName,
                suspendDuration,
                TkEndpointManager.SuspendReason.RemoteDailyLimit,
                notifyOtherNodes: false,
                emitNotification: false);

            var endpointConfigs = await TkEndpointCore.GetEndpointsByAccountReadonlyAsync(
                account.id,
                parseEndpoints: account.parseEndpoint);
            int endpointCounterId = endpointConfigs
                .FirstOrDefault(endpoint => endpoint.ep_id == EndpointId)
                ?.id ?? EndpointId;

            int requestCount = await RiskControlCore.GetAllNodesTkEndpointCallsAsync(
                account.id,
                endpointCounterId,
                triggerTime.ToString("yyyyMMdd"));
            int successCount = await TkLogCore.GetTotalAsync(
                $":parse_total:{TkChannelEnum.tb}_{account.id}:success:{triggerTime:yyyyMMdd}");

            bool shouldBroadcast = false;
            try
            {
                using var conn = CenterHub.GetOpenConnection();
                const string sql = @"
INSERT IGNORE INTO tk_aiopen_risk_daily
    (report_date, account_id, account_name, endpoint_id, endpoint_name,
     reason_code, reason, biz_error_code, result_code, first_trigger_time,
     request_count_at_trigger, success_count_at_trigger, release_time,
     create_time)
VALUES
    (@report_date, @account_id, @account_name, @endpoint_id, @endpoint_name,
     @reason_code, LEFT(@reason, 512), @biz_error_code, @result_code, @first_trigger_time,
     @request_count_at_trigger, @success_count_at_trigger, @release_time,
     @create_time);";

                int inserted = await conn.ExecuteAsync(sql, new
                {
                    report_date = triggerTime.Date,
                    account_id = account.id,
                    account_name = account.company,
                    endpoint_id = EndpointId,
                    endpoint_name = EndpointName,
                    reason_code = DailyLimitReasonCode,
                    reason,
                    biz_error_code = bizErrorCode,
                    result_code = resultCode,
                    first_trigger_time = triggerTime,
                    request_count_at_trigger = requestCount,
                    success_count_at_trigger = successCount,
                    release_time = releaseTime,
                    create_time = triggerTime
                });
                shouldBroadcast = inserted > 0;
            }
            catch (Exception ex)
            {
                // Suspending all nodes is more important than report persistence. Broadcast
                // even when the migration has not yet been applied or DB is down.
                shouldBroadcast = true;
                _ = new LoggerLibrary("AiOpenRiskControl", "save_error")
                    .Info($"accountId={account.id}, requestCount={requestCount}, reason={reason}")
                    .Info(ex.Message, ex.StackTrace)
                    .SaveAsync();
            }

            if (shouldBroadcast)
            {
                await TkEndpointManager.SuspendForDurationAsync(
                    account.id,
                    EndpointName,
                    suspendDuration,
                    TkEndpointManager.SuspendReason.RemoteDailyLimit,
                    notifyOtherNodes: true,
                    emitNotification: true);
            }
        }

        public static async Task<Dictionary<int, AiOpenRiskDailyDTO>> GetTodayEventsAsync(IEnumerable<int> accountIds)
        {
            int[] ids = accountIds.Where(id => id > 0).Distinct().ToArray();
            if (ids.Length == 0)
            {
                return [];
            }

            try
            {
                using var conn = CenterHub.GetOpenConnection();
                const string sql = @"
SELECT id, report_date, account_id, account_name, endpoint_id, endpoint_name,
       reason_code, reason, biz_error_code, result_code, first_trigger_time,
       request_count_at_trigger, success_count_at_trigger, release_time,
       create_time
FROM tk_aiopen_risk_daily
WHERE report_date = @report_date
  AND account_id IN @account_ids
  AND reason_code = @reason_code;";

                var rows = await conn.QueryAsync<AiOpenRiskDailyDTO>(sql, new
                {
                    report_date = DateTime.Now.Date,
                    account_ids = ids,
                    reason_code = DailyLimitReasonCode
                });

                return rows.ToDictionary(row => row.account_id);
            }
            catch (Exception ex)
            {
                _ = new LoggerLibrary("AiOpenRiskControl", "query_error")
                    .Info(ex.Message, ex.StackTrace)
                    .SaveAsync();
                return [];
            }
        }
    }
}