ai-chatbot/lib/canvas/image-concurrency.ts

import "server-only";

import { getSystemSetting } from "@/lib/db/queries";
import {
  CANVAS_IMAGE_TASK_QUEUE_CONFIG_KEY,
  DEFAULT_CANVAS_IMAGE_BATCH_DISPATCH_INTERVAL_MS,
  DEFAULT_CANVAS_IMAGE_TASK_CONCURRENCY,
  type CanvasImageTaskQueueConfig,
  normalizeCanvasImageTaskQueueConfig,
} from "./image-task-config";
import {
  estimateImageDispatchWaitMs,
  getAdaptiveImageQueueOwnerLimit,
} from "./image-queue-fairness";

const CONCURRENCY_CONFIG_CACHE_MS = 5 * 1000;
const RECENT_EXECUTION_SAMPLE_LIMIT = 24;
const DEFAULT_ESTIMATED_WAIT_MS = 45 * 1000;
const concurrencyStateKey = "__canvas_image_model_concurrency__";

interface ModelWaiter {
  ticketId: string;
  ownerKey: string;
  resolve: () => void;
  onQueued?: (status: CanvasImageModelQueueStatus) => void | Promise<void>;
  lastStatusSignature?: string;
  debugLabel?: string;
  debugMeta?: Record<string, unknown>;
}

interface ModelConcurrencyState {
  running: Map<string, string>;
  waiters: ModelWaiter[];
  dispatchTimer: ReturnType<typeof setTimeout> | null;
  delayedUntil: number | null;
  recentExecutionDurationsMs: number[];
}

interface CanvasImageConcurrencyState {
  models: Map<string, ModelConcurrencyState>;
}

let cachedConfig:
  | {
      expiresAt: number;
      value: CanvasImageTaskQueueConfig;
    }
  | null = null;

function getGlobalConcurrencyState(): CanvasImageConcurrencyState {
  const globalState = globalThis as typeof globalThis & {
    [concurrencyStateKey]?: CanvasImageConcurrencyState;
  };

  if (!globalState[concurrencyStateKey]) {
    globalState[concurrencyStateKey] = {
      models: new Map<string, ModelConcurrencyState>(),
    };
  }

  return globalState[concurrencyStateKey] as CanvasImageConcurrencyState;
}

function getOrCreateModelState(modelKey: string) {
  const state = getGlobalConcurrencyState();
  const existing = state.models.get(modelKey);
  if (existing) {
    return existing;
  }

  const created: ModelConcurrencyState = {
    running: new Map<string, string>(),
    waiters: [],
    dispatchTimer: null,
    delayedUntil: null,
    recentExecutionDurationsMs: [],
  };
  state.models.set(modelKey, created);
  return created;
}

function clearModelDispatchTimer(modelState: ModelConcurrencyState) {
  if (modelState.dispatchTimer) {
    clearTimeout(modelState.dispatchTimer);
    modelState.dispatchTimer = null;
  }

  modelState.delayedUntil = null;
}

function cleanupModelState(modelKey: string) {
  const state = getGlobalConcurrencyState();
  const modelState = state.models.get(modelKey);
  if (!modelState) {
    return;
  }

  if (
    modelState.running.size === 0 &&
    modelState.waiters.length === 0 &&
    modelState.recentExecutionDurationsMs.length === 0 &&
    !modelState.dispatchTimer &&
    (!modelState.delayedUntil || modelState.delayedUntil <= Date.now())
  ) {
    clearModelDispatchTimer(modelState);
    state.models.delete(modelKey);
  }
}

export function normalizeCanvasImageModelKey(value: unknown) {
  if (typeof value !== "string") {
    return "__default__";
  }

  const normalized = value.trim();
  return normalized || "__default__";
}

function normalizeCanvasImageOwnerKey(value: unknown) {
  if (typeof value !== "string") {
    return "__anonymous__";
  }

  const normalized = value.trim();
  return normalized || "__anonymous__";
}

function getRunningCountForOwner(
  modelState: ModelConcurrencyState,
  ownerKey: string
) {
  let count = 0;
  for (const runningOwnerKey of modelState.running.values()) {
    if (runningOwnerKey === ownerKey) {
      count += 1;
    }
  }
  return count;
}

function hasCompetingWaiter(
  modelState: ModelConcurrencyState,
  ownerKey: string
) {
  return modelState.waiters.some((waiter) => waiter.ownerKey !== ownerKey);
}

export async function getCanvasImageTaskQueueConfig(): Promise<CanvasImageTaskQueueConfig> {
  const now = Date.now();
  if (cachedConfig && cachedConfig.expiresAt > now) {
    return cachedConfig.value;
  }

  const envConcurrency = process.env.CANVAS_IMAGE_TASK_CONCURRENCY;
  const envBatchDispatchIntervalMs = process.env.CANVAS_IMAGE_TASK_BATCH_DISPATCH_INTERVAL_MS;
  let config = normalizeCanvasImageTaskQueueConfig(undefined, {
    envConcurrency,
    envBatchDispatchIntervalMs,
  });

  try {
    const setting = await getSystemSetting({
      key: CANVAS_IMAGE_TASK_QUEUE_CONFIG_KEY,
    });
    if (setting?.value !== undefined && setting.value !== null) {
      config = normalizeCanvasImageTaskQueueConfig(setting.value, {
        envConcurrency,
        envBatchDispatchIntervalMs,
      });
    }
  } catch (error) {
    console.warn("[Canvas Image Concurrency] Failed to load DB config, using fallback", {
      error: error instanceof Error ? error.message : String(error),
    });
  }

  cachedConfig = {
    expiresAt: now + CONCURRENCY_CONFIG_CACHE_MS,
    value: {
      concurrency: config.concurrency || DEFAULT_CANVAS_IMAGE_TASK_CONCURRENCY,
      batchDispatchIntervalMs: Number.isFinite(config.batchDispatchIntervalMs)
        ? config.batchDispatchIntervalMs
        : DEFAULT_CANVAS_IMAGE_BATCH_DISPATCH_INTERVAL_MS,
    },
  };

  return cachedConfig.value;
}

export async function getCanvasImageSingleModelConcurrency(): Promise<number> {
  return (await getCanvasImageTaskQueueConfig()).concurrency;
}

export async function getCanvasImageBatchDispatchIntervalMs(): Promise<number> {
  return (await getCanvasImageTaskQueueConfig()).batchDispatchIntervalMs;
}

export function getCanvasImageRunningCountForModel(modelKey: string) {
  const normalizedModelKey = normalizeCanvasImageModelKey(modelKey);
  return getGlobalConcurrencyState().models.get(normalizedModelKey)?.running.size || 0;
}

export type CanvasImageModelQueueStatus = {
  status: "pending" | "running" | "succeeded" | "failed";
  queuePosition: number;
  sameModelPendingCount: number;
  sameModelRunningCount: number;
  concurrencyLimit: number;
  dispatchIntervalMs: number;
  estimatedWaitMs: number;
  source: "model_queue";
};

function getRecentAverageExecutionMs(modelState: ModelConcurrencyState) {
  const samples = modelState.recentExecutionDurationsMs.filter(
    (value) => Number.isFinite(value) && value > 0
  );
  if (samples.length === 0) {
    return DEFAULT_ESTIMATED_WAIT_MS;
  }
  const total = samples.reduce((sum, value) => sum + value, 0);
  return Math.max(1000, Math.round(total / samples.length));
}

function buildModelQueueSnapshot(params: {
  modelState: ModelConcurrencyState;
  queuePosition: number;
  sameModelPendingCount: number;
  sameModelRunningCount: number;
  concurrencyLimit: number;
  dispatchIntervalMs: number;
}): CanvasImageModelQueueStatus {
  const averageExecutionMs = getRecentAverageExecutionMs(params.modelState);
  const delayRemainingMs = Math.max(
    0,
    Number(params.modelState.delayedUntil || 0) - Date.now()
  );
  const estimatedWaitMs =
    params.sameModelPendingCount <= 0
      ? 0
      : params.dispatchIntervalMs > 0
        ? estimateImageDispatchWaitMs({
            queuePosition: params.queuePosition,
            dispatchIntervalMs: params.dispatchIntervalMs,
            delayRemainingMs,
          })
        : Math.max(
            averageExecutionMs,
            Math.ceil(
              (params.queuePosition + 1) / Math.max(1, params.concurrencyLimit)
            ) * averageExecutionMs
          );

  return {
    status: "pending",
    queuePosition: Math.max(0, params.queuePosition),
    sameModelPendingCount: Math.max(0, params.sameModelPendingCount),
    sameModelRunningCount: Math.max(0, params.sameModelRunningCount),
    concurrencyLimit: Math.max(1, params.concurrencyLimit),
    dispatchIntervalMs: Math.max(0, params.dispatchIntervalMs),
    estimatedWaitMs,
    source: "model_queue",
  };
}

function notifyModelWaiters(
  modelKey: string,
  modelState: ModelConcurrencyState,
  concurrencyLimit: number,
  dispatchIntervalMs: number
) {
  const sameModelPendingCount = modelState.waiters.length;
  for (const [queuePosition, waiter] of modelState.waiters.entries()) {
    if (typeof waiter.onQueued !== "function") {
      continue;
    }

    const queueStatus = buildModelQueueSnapshot({
      modelState,
      queuePosition,
      sameModelPendingCount,
      sameModelRunningCount: modelState.running.size,
      concurrencyLimit,
      dispatchIntervalMs,
    });
    const statusSignature = [
      queueStatus.queuePosition,
      queueStatus.sameModelPendingCount,
      queueStatus.sameModelRunningCount,
      queueStatus.concurrencyLimit,
      queueStatus.dispatchIntervalMs,
      queueStatus.estimatedWaitMs,
    ].join(":");
    if (waiter.lastStatusSignature === statusSignature) {
      continue;
    }

    waiter.lastStatusSignature = statusSignature;
    Promise.resolve(waiter.onQueued(queueStatus)).catch((error) => {
      console.warn("[Canvas Image Concurrency] onQueued callback failed:", {
        debugLabel: waiter.debugLabel || "canvas-image",
        modelKey,
        error: error instanceof Error ? error.message : String(error),
        ...waiter.debugMeta,
      });
    });
  }
}

export async function getCanvasImageModelQueueStatus(params: {
  modelKey: string;
}): Promise<CanvasImageModelQueueStatus> {
  const normalizedModelKey = normalizeCanvasImageModelKey(params.modelKey);
  const { concurrency: maxConcurrent, batchDispatchIntervalMs } =
    await getCanvasImageTaskQueueConfig();
  const modelState =
    getGlobalConcurrencyState().models.get(normalizedModelKey) || {
      running: new Map<string, string>(),
      waiters: [],
      dispatchTimer: null,
      delayedUntil: null,
      recentExecutionDurationsMs: [],
    };
  return buildModelQueueSnapshot({
    modelState,
    queuePosition: modelState.waiters.length > 0 ? modelState.waiters.length - 1 : 0,
    sameModelPendingCount: modelState.waiters.length,
    sameModelRunningCount: modelState.running.size,
    concurrencyLimit: maxConcurrent,
    dispatchIntervalMs: batchDispatchIntervalMs,
  });
}

export function recordCanvasImageModelExecutionDuration(
  modelKey: string,
  durationMs: number
) {
  if (!Number.isFinite(durationMs) || durationMs <= 0) {
    return;
  }

  const normalizedModelKey = normalizeCanvasImageModelKey(modelKey);
  const modelState = getOrCreateModelState(normalizedModelKey);
  modelState.recentExecutionDurationsMs.push(durationMs);
  if (modelState.recentExecutionDurationsMs.length > RECENT_EXECUTION_SAMPLE_LIMIT) {
    modelState.recentExecutionDurationsMs.splice(
      0,
      modelState.recentExecutionDurationsMs.length - RECENT_EXECUTION_SAMPLE_LIMIT
    );
  }
}

async function drainModelWaiters(
  modelKey: string,
  _trigger: "arrival" | "slot-release" | "timer"
) {
  const normalizedModelKey = normalizeCanvasImageModelKey(modelKey);
  const modelState = getGlobalConcurrencyState().models.get(normalizedModelKey);
  if (!modelState || modelState.waiters.length === 0) {
    cleanupModelState(normalizedModelKey);
    return;
  }

  const { concurrency: maxConcurrent, batchDispatchIntervalMs } =
    await getCanvasImageTaskQueueConfig();
  notifyModelWaiters(
    normalizedModelKey,
    modelState,
    maxConcurrent,
    batchDispatchIntervalMs
  );
  if (modelState.running.size >= maxConcurrent) {
    return;
  }

  const now = Date.now();
  if (modelState.delayedUntil && modelState.delayedUntil > now) {
    return;
  }

  clearModelDispatchTimer(modelState);
  while (modelState.running.size < maxConcurrent && modelState.waiters.length > 0) {
    const waiterIndex = modelState.waiters.findIndex(
      (item) => {
        const ownerLimit = getAdaptiveImageQueueOwnerLimit({
          maxConcurrent,
          hasCompetingOwner: hasCompetingWaiter(modelState, item.ownerKey),
        });
        return getRunningCountForOwner(modelState, item.ownerKey) < ownerLimit;
      }
    );
    if (waiterIndex < 0) {
      break;
    }

    const [waiter] = modelState.waiters.splice(waiterIndex, 1);
    if (!waiter) {
      break;
    }
    modelState.running.set(waiter.ticketId, waiter.ownerKey);
    waiter.resolve();
    if (batchDispatchIntervalMs > 0) {
      modelState.delayedUntil = Date.now() + batchDispatchIntervalMs;
      modelState.dispatchTimer = setTimeout(() => {
        const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
        if (!state) return;
        clearModelDispatchTimer(state);
        void drainModelWaiters(normalizedModelKey, "timer");
      }, batchDispatchIntervalMs);
      break;
    }
  }

  notifyModelWaiters(
    normalizedModelKey,
    modelState,
    maxConcurrent,
    batchDispatchIntervalMs
  );

  cleanupModelState(normalizedModelKey);
}

export async function acquireCanvasImageModelSlot(params: {
  modelKey: string;
  ownerKey?: string;
  debugLabel?: string;
  debugMeta?: Record<string, unknown>;
  onQueued?: (status: CanvasImageModelQueueStatus) => void | Promise<void>;
}) {
  const normalizedModelKey = normalizeCanvasImageModelKey(params.modelKey);
  const normalizedOwnerKey = normalizeCanvasImageOwnerKey(params.ownerKey);
  const {
    concurrency: maxConcurrent,
    batchDispatchIntervalMs,
  } = await getCanvasImageTaskQueueConfig();
  const modelState = getOrCreateModelState(normalizedModelKey);
  const ticketId = `${normalizedModelKey}:${Date.now()}:${Math.random().toString(36).slice(2, 10)}`;
  const waitingSince = Date.now();
  const hasActiveDelay =
    modelState.delayedUntil !== null && modelState.delayedUntil > waitingSince;
  if (!hasActiveDelay && modelState.dispatchTimer) {
    clearModelDispatchTimer(modelState);
  }
  const ownerLimit = getAdaptiveImageQueueOwnerLimit({
    maxConcurrent,
    hasCompetingOwner: hasCompetingWaiter(modelState, normalizedOwnerKey),
  });

  if (
    modelState.waiters.length === 0 &&
    !hasActiveDelay &&
    modelState.running.size < maxConcurrent &&
    getRunningCountForOwner(modelState, normalizedOwnerKey) < ownerLimit
  ) {
    modelState.running.set(ticketId, normalizedOwnerKey);
    if (batchDispatchIntervalMs > 0) {
      modelState.delayedUntil = Date.now() + batchDispatchIntervalMs;
      modelState.dispatchTimer = setTimeout(() => {
        const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
        if (!state) return;
        clearModelDispatchTimer(state);
        void drainModelWaiters(normalizedModelKey, "timer");
      }, batchDispatchIntervalMs);
    }
    return {
      release: async () => {
        const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
        if (!state) return;
        state.running.delete(ticketId);
        await drainModelWaiters(normalizedModelKey, "slot-release");
      },
      waitMs: 0,
      maxConcurrent,
      batchDispatchIntervalMs,
    };
  }

  await new Promise<void>((resolve) => {
    modelState.waiters.push({
      ticketId,
      ownerKey: normalizedOwnerKey,
      resolve,
      onQueued: params.onQueued,
      debugLabel: params.debugLabel,
      debugMeta: params.debugMeta,
    });
    notifyModelWaiters(
      normalizedModelKey,
      modelState,
      maxConcurrent,
      batchDispatchIntervalMs
    );
    void drainModelWaiters(normalizedModelKey, "arrival");
  });

  const waitMs = Date.now() - waitingSince;
  if (waitMs > 250) {
    console.info("[Canvas Image Concurrency] Waited for model slot:", {
      debugLabel: params.debugLabel || "canvas-image",
      modelKey: normalizedModelKey,
      waitMs,
      batchDispatchIntervalMs,
      ...params.debugMeta,
    });
  }

  return {
    release: async () => {
      const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
      if (!state) return;
      state.running.delete(ticketId);
      await drainModelWaiters(normalizedModelKey, "slot-release");
    },
    waitMs,
    maxConcurrent,
    batchDispatchIntervalMs,
  };
}