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_TASK_CONCURRENCY,
  normalizeCanvasImageTaskQueueConfig,
} from "./image-task-config";

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

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

interface ModelConcurrencyState {
  running: Set<string>;
  waiters: ModelWaiter[];
}

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

let cachedConcurrency:
  | {
      expiresAt: number;
      value: number;
    }
  | 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 Set<string>(),
    waiters: [],
  };
  state.models.set(modelKey, created);
  return created;
}

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) {
    state.models.delete(modelKey);
  }
}

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

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

export async function getCanvasImageSingleModelConcurrency(): Promise<number> {
  const now = Date.now();
  if (cachedConcurrency && cachedConcurrency.expiresAt > now) {
    return cachedConcurrency.value;
  }

  const envConcurrency = process.env.CANVAS_IMAGE_TASK_CONCURRENCY;
  let value = normalizeCanvasImageTaskQueueConfig(undefined, {
    envConcurrency,
  }).concurrency;

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

  cachedConcurrency = {
    expiresAt: now + CONCURRENCY_CONFIG_CACHE_MS,
    value: value || DEFAULT_CANVAS_IMAGE_TASK_CONCURRENCY,
  };

  return cachedConcurrency.value;
}

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

async function drainModelWaiters(modelKey: string) {
  const normalizedModelKey = normalizeCanvasImageModelKey(modelKey);
  const modelState = getGlobalConcurrencyState().models.get(normalizedModelKey);
  if (!modelState || modelState.waiters.length === 0) {
    cleanupModelState(normalizedModelKey);
    return;
  }

  const maxConcurrent = await getCanvasImageSingleModelConcurrency();
  while (modelState.running.size < maxConcurrent && modelState.waiters.length > 0) {
    const waiter = modelState.waiters.shift();
    if (!waiter) {
      break;
    }
    modelState.running.add(waiter.ticketId);
    waiter.resolve();
  }

  cleanupModelState(normalizedModelKey);
}

export async function acquireCanvasImageModelSlot(params: {
  modelKey: string;
  debugLabel?: string;
  debugMeta?: Record<string, unknown>;
}) {
  const normalizedModelKey = normalizeCanvasImageModelKey(params.modelKey);
  const maxConcurrent = await getCanvasImageSingleModelConcurrency();
  const modelState = getOrCreateModelState(normalizedModelKey);
  const ticketId = `${normalizedModelKey}:${Date.now()}:${Math.random().toString(36).slice(2, 10)}`;
  const waitingSince = Date.now();

  if (modelState.waiters.length === 0 && modelState.running.size < maxConcurrent) {
    modelState.running.add(ticketId);
    return {
      release: async () => {
        const state = getGlobalConcurrencyState().models.get(normalizedModelKey);
        if (!state) return;
        state.running.delete(ticketId);
        await drainModelWaiters(normalizedModelKey);
      },
      waitMs: 0,
      maxConcurrent,
    };
  }

  await new Promise<void>((resolve) => {
    modelState.waiters.push({
      ticketId,
      resolve,
    });
    void drainModelWaiters(normalizedModelKey);
  });

  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,
      ...params.debugMeta,
    });
  }

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