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";

const CONCURRENCY_CONFIG_CACHE_MS = 5 * 1000;
const concurrencyStateKey = "__canvas_image_model_concurrency__";

interface ModelWaiter {
  ticketId: string;
  ownerKey: string;
  resolve: () => void;
}

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

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,
  };
  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) {
    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 getMaxConcurrentPerOwner(maxConcurrent: number) {
  return Math.max(1, Math.ceil(Math.max(1, maxConcurrent) / 2));
}

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

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;
}

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();
  if (modelState.running.size >= maxConcurrent) {
    return;
  }

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

  if (trigger === "slot-release" && batchDispatchIntervalMs > 0) {
    if (!modelState.dispatchTimer) {
      modelState.delayedUntil = now + batchDispatchIntervalMs;
      modelState.dispatchTimer = setTimeout(() => {
        const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
        if (!state) return;
        clearModelDispatchTimer(state);
        void drainModelWaiters(normalizedModelKey, "timer");
      }, batchDispatchIntervalMs);
    }
    return;
  }

  clearModelDispatchTimer(modelState);
  while (modelState.running.size < maxConcurrent && modelState.waiters.length > 0) {
    const maxConcurrentPerOwner = getMaxConcurrentPerOwner(maxConcurrent);
    const waiterIndex = modelState.waiters.findIndex(
      (item) =>
        getRunningCountForOwner(modelState, item.ownerKey) < maxConcurrentPerOwner
    );
    if (waiterIndex < 0) {
      break;
    }

    const [waiter] = modelState.waiters.splice(waiterIndex, 1);
    if (!waiter) {
      break;
    }
    modelState.running.set(waiter.ticketId, waiter.ownerKey);
    waiter.resolve();
  }

  cleanupModelState(normalizedModelKey);
}

export async function acquireCanvasImageModelSlot(params: {
  modelKey: string;
  ownerKey?: string;
  debugLabel?: string;
  debugMeta?: Record<string, unknown>;
}) {
  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;
  const maxConcurrentPerOwner = getMaxConcurrentPerOwner(maxConcurrent);

  if (
    modelState.waiters.length === 0 &&
    !hasActiveDelay &&
    modelState.running.size < maxConcurrent &&
    getRunningCountForOwner(modelState, normalizedOwnerKey) < maxConcurrentPerOwner
  ) {
    modelState.running.set(ticketId, normalizedOwnerKey);
    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,
    });
    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,
  };
}