ai-chatbot/lib/canvas/image-task-queue.ts

import "server-only";

import { and, asc, eq, lt, notInArray, sql } from "drizzle-orm";
import { db } from "@/lib/db";
import { canvasImageTask } from "@/lib/db/schema";
import {
  getCanvasImageBatchDispatchIntervalMs,
  getCanvasImageSingleModelConcurrency,
  normalizeCanvasImageModelKey,
} from "./image-concurrency";
import {
  generateCanvasImage,
  type CanvasGenerateImageRequest,
  type CanvasGenerateImageResult,
} from "./generate-image";

type CanvasImageTaskStatus = "pending" | "running" | "succeeded" | "failed";

interface PendingTaskEntry {
  taskId: string;
  modelKey: string;
  userId: string;
}

interface RunningTaskEntry {
  modelKey: string;
  userId: string;
}

interface QueueState {
  pending: PendingTaskEntry[];
  pendingSet: Set<string>;
  running: Map<string, RunningTaskEntry>;
  runningByModel: Map<string, Set<string>>;
  runningByModelUser: Map<string, Set<string>>;
  nextBatchEligibleAt: Map<string, number>;
  nextBatchTimers: Map<string, ReturnType<typeof setTimeout>>;
  awaitingNextBatch: Set<string>;
  draining: boolean;
}

const STALE_RUNNING_MS = 5 * 60 * 1000;
const PENDING_LOOKAHEAD_LIMIT = 100;
const queueStateKey = "__canvas_image_task_queue_state__";

function getQueueState(): QueueState {
  const globalObj = globalThis as typeof globalThis & {
    [queueStateKey]?: QueueState;
  };
  if (!globalObj[queueStateKey]) {
    globalObj[queueStateKey] = {
      pending: [],
      pendingSet: new Set<string>(),
      running: new Map<string, RunningTaskEntry>(),
      runningByModel: new Map<string, Set<string>>(),
      runningByModelUser: new Map<string, Set<string>>(),
      nextBatchEligibleAt: new Map<string, number>(),
      nextBatchTimers: new Map<string, ReturnType<typeof setTimeout>>(),
      awaitingNextBatch: new Set<string>(),
      draining: false,
    };
  }
  return globalObj[queueStateKey] as QueueState;
}

function normalizeTaskModelKey(request: unknown) {
  if (!request || typeof request !== "object" || Array.isArray(request)) {
    return "__default__";
  }

  const rawModel = (request as Record<string, unknown>).model;
  return normalizeCanvasImageModelKey(rawModel);
}

function getModelUserKey(modelKey: string, userId: string) {
  return `${modelKey}::${userId || "__anonymous__"}`;
}

function getMaxConcurrentPerUser(maxConcurrentPerModel: number) {
  return Math.max(1, Math.ceil(Math.max(1, maxConcurrentPerModel) / 2));
}

async function getTaskQueueMeta(taskId: string) {
  const rows = await db
    .select({
      request: canvasImageTask.request,
      userId: canvasImageTask.userId,
    })
    .from(canvasImageTask)
    .where(eq(canvasImageTask.id, taskId))
    .limit(1);

  const row = rows[0];
  return {
    modelKey: normalizeTaskModelKey(row?.request),
    userId: String(row?.userId || "__anonymous__"),
  };
}

function getRunningCountForModel(modelKey: string) {
  const state = getQueueState();
  return state.runningByModel.get(modelKey)?.size || 0;
}

function getRunningCountForModelUser(modelKey: string, userId: string) {
  const state = getQueueState();
  return state.runningByModelUser.get(getModelUserKey(modelKey, userId))?.size || 0;
}

function markRunningTask(taskId: string, modelKey: string, userId: string) {
  const state = getQueueState();
  state.running.set(taskId, { modelKey, userId });
  const tasks = state.runningByModel.get(modelKey) || new Set<string>();
  tasks.add(taskId);
  state.runningByModel.set(modelKey, tasks);

  const modelUserKey = getModelUserKey(modelKey, userId);
  const userTasks = state.runningByModelUser.get(modelUserKey) || new Set<string>();
  userTasks.add(taskId);
  state.runningByModelUser.set(modelUserKey, userTasks);
}

function releaseRunningTask(taskId: string) {
  const state = getQueueState();
  const entry = state.running.get(taskId);
  if (!entry) {
    return null;
  }

  state.running.delete(taskId);
  const { modelKey, userId } = entry;
  const tasks = state.runningByModel.get(modelKey);
  if (tasks) {
    tasks.delete(taskId);
    if (tasks.size === 0) {
      state.runningByModel.delete(modelKey);
    }
  }

  const modelUserKey = getModelUserKey(modelKey, userId);
  const userTasks = state.runningByModelUser.get(modelUserKey);
  if (userTasks) {
    userTasks.delete(taskId);
    if (userTasks.size === 0) {
      state.runningByModelUser.delete(modelUserKey);
    }
  }

  return modelKey;
}

function isModelBatchDispatchDelayed(modelKey: string) {
  const state = getQueueState();
  const delayedUntil = state.nextBatchEligibleAt.get(modelKey);
  if (!delayedUntil) {
    return false;
  }

  if (delayedUntil <= Date.now()) {
    state.nextBatchEligibleAt.delete(modelKey);
    const timer = state.nextBatchTimers.get(modelKey);
    if (timer) {
      clearTimeout(timer);
      state.nextBatchTimers.delete(modelKey);
    }
    return false;
  }

  return true;
}

function markModelAwaitingNextBatch(modelKey: string) {
  getQueueState().awaitingNextBatch.add(modelKey);
}

function scheduleModelNextBatch(modelKey: string, delayMs: number) {
  const state = getQueueState();
  if (delayMs <= 0) {
    state.awaitingNextBatch.delete(modelKey);
    return false;
  }

  if (state.nextBatchTimers.has(modelKey)) {
    return true;
  }

  state.awaitingNextBatch.delete(modelKey);
  state.nextBatchEligibleAt.set(modelKey, Date.now() + delayMs);
  const timer = setTimeout(() => {
    const latestState = getQueueState();
    latestState.nextBatchTimers.delete(modelKey);
    latestState.nextBatchEligibleAt.delete(modelKey);
    void drainCanvasImageTaskQueue();
  }, delayMs);
  state.nextBatchTimers.set(modelKey, timer);
  return true;
}

async function enqueueLocalTask(
  taskId: string,
  meta?: { modelKey?: string; userId?: string } | string
) {
  const state = getQueueState();
  if (!taskId || state.running.has(taskId) || state.pendingSet.has(taskId)) {
    return;
  }

  const resolvedMeta =
    typeof meta === "string"
      ? { modelKey: meta, userId: undefined }
      : meta || {};
  const dbMeta =
    resolvedMeta.modelKey && resolvedMeta.userId
      ? null
      : await getTaskQueueMeta(taskId);
  const resolvedModelKey = resolvedMeta.modelKey || dbMeta?.modelKey || "__default__";
  const resolvedUserId = resolvedMeta.userId || dbMeta?.userId || "__anonymous__";
  state.pending.push({
    taskId,
    modelKey: resolvedModelKey,
    userId: resolvedUserId,
  });
  state.pendingSet.add(taskId);

  const maxConcurrentPerModel = await getCanvasImageSingleModelConcurrency();
  if (
    getRunningCountForModel(resolvedModelKey) >= maxConcurrentPerModel ||
    getRunningCountForModelUser(resolvedModelKey, resolvedUserId) >=
      getMaxConcurrentPerUser(maxConcurrentPerModel)
  ) {
    markModelAwaitingNextBatch(resolvedModelKey);
  }
}

function hasModelCapacity(
  modelKey: string,
  userId: string,
  maxConcurrentPerModel: number
) {
  return (
    getRunningCountForModel(modelKey) < maxConcurrentPerModel &&
    getRunningCountForModelUser(modelKey, userId) <
      getMaxConcurrentPerUser(maxConcurrentPerModel)
  );
}

async function claimNextPendingTask(
  maxConcurrentPerModel: number
): Promise<PendingTaskEntry | null> {
  const state = getQueueState();
  for (let index = 0; index < state.pending.length; index += 1) {
    const entry = state.pending[index];
    if (!entry) continue;
    if (state.running.has(entry.taskId)) continue;
    if (isModelBatchDispatchDelayed(entry.modelKey)) continue;
    if (!hasModelCapacity(entry.modelKey, entry.userId, maxConcurrentPerModel)) {
      markModelAwaitingNextBatch(entry.modelKey);
      continue;
    }

    state.pending.splice(index, 1);
    state.pendingSet.delete(entry.taskId);
    return entry;
  }

  const runningIds = Array.from(state.running.keys());
  const rows = await db
    .select({
      id: canvasImageTask.id,
      request: canvasImageTask.request,
      userId: canvasImageTask.userId,
    })
    .from(canvasImageTask)
    .where(
      runningIds.length > 0
        ? and(eq(canvasImageTask.status, "pending"), notInArray(canvasImageTask.id, runningIds))
        : eq(canvasImageTask.status, "pending")
    )
    .orderBy(asc(canvasImageTask.createdAt))
    .limit(PENDING_LOOKAHEAD_LIMIT);

  for (const row of rows) {
    if (!row?.id) continue;
    if (state.pendingSet.has(row.id) || state.running.has(row.id)) continue;

    const modelKey = normalizeTaskModelKey(row.request);
    if (isModelBatchDispatchDelayed(modelKey)) {
      continue;
    }
    const userId = String(row.userId || "__anonymous__");
    if (!hasModelCapacity(modelKey, userId, maxConcurrentPerModel)) {
      markModelAwaitingNextBatch(modelKey);
      continue;
    }

    return {
      taskId: row.id,
      modelKey,
      userId,
    };
  }

  return null;
}

async function runSingleTask(entry: PendingTaskEntry) {
  try {
    const [task] = await db
      .update(canvasImageTask)
      .set({
        status: "running",
        startedAt: new Date(),
        updatedAt: new Date(),
        completedAt: null,
        error: null,
        attempts: sql`${canvasImageTask.attempts} + 1`,
      })
      .where(and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "pending")))
      .returning({
        id: canvasImageTask.id,
        request: canvasImageTask.request,
        userId: canvasImageTask.userId,
      });

    if (!task) {
      return;
    }

    const result = await generateCanvasImage({
      ...(task.request as CanvasGenerateImageRequest),
      ownerKey: task.userId,
    });
    await db
      .update(canvasImageTask)
      .set({
        status: "succeeded",
        result: result as CanvasGenerateImageResult,
        error: null,
        completedAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(canvasImageTask.id, entry.taskId));
  } catch (error) {
    await db
      .update(canvasImageTask)
      .set({
        status: "failed",
        error: error instanceof Error ? error.message : String(error),
        completedAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(canvasImageTask.id, entry.taskId));
  } finally {
    const modelKey = releaseRunningTask(entry.taskId);
    if (modelKey) {
      const batchDispatchIntervalMs = await getCanvasImageBatchDispatchIntervalMs();
      if (getQueueState().awaitingNextBatch.has(modelKey)) {
        scheduleModelNextBatch(modelKey, batchDispatchIntervalMs);
      }
    }
    void drainCanvasImageTaskQueue();
  }
}

export async function drainCanvasImageTaskQueue() {
  const state = getQueueState();
  if (state.draining) return;

  state.draining = true;
  try {
    const maxConcurrentPerModel = await getCanvasImageSingleModelConcurrency();
    while (true) {
      const nextTask = await claimNextPendingTask(maxConcurrentPerModel);
      if (!nextTask) break;
      if (state.running.has(nextTask.taskId)) continue;
      markRunningTask(nextTask.taskId, nextTask.modelKey, nextTask.userId);
      void runSingleTask(nextTask);
    }
  } finally {
    state.draining = false;
  }
}

export async function enqueueCanvasImageTask(taskId: string) {
  await enqueueLocalTask(taskId);
  await drainCanvasImageTaskQueue();
}

export async function createCanvasImageTask(params: {
  userId: string;
  request: CanvasGenerateImageRequest;
}) {
  const [task] = await db
    .insert(canvasImageTask)
      .values({
        userId: params.userId,
        status: "pending",
        request: params.request as unknown as Record<string, unknown>,
        attempts: 0,
        createdAt: new Date(),
        updatedAt: new Date(),
      })
    .returning({
      id: canvasImageTask.id,
      status: canvasImageTask.status,
      createdAt: canvasImageTask.createdAt,
    });

  return task;
}

export async function getCanvasImageTaskForUser(params: {
  taskId: string;
  userId: string;
}) {
  const rows = await db
    .select({
      id: canvasImageTask.id,
      status: canvasImageTask.status,
      request: canvasImageTask.request,
      result: canvasImageTask.result,
      error: canvasImageTask.error,
      attempts: canvasImageTask.attempts,
      startedAt: canvasImageTask.startedAt,
      completedAt: canvasImageTask.completedAt,
      createdAt: canvasImageTask.createdAt,
      updatedAt: canvasImageTask.updatedAt,
    })
    .from(canvasImageTask)
    .where(and(eq(canvasImageTask.id, params.taskId), eq(canvasImageTask.userId, params.userId)))
    .limit(1);

  const task = rows[0];
  if (!task) return null;

  if (task.status === "running" && task.startedAt) {
    const staleBefore = new Date(Date.now() - STALE_RUNNING_MS);
    if (task.startedAt < staleBefore) {
      await db
        .update(canvasImageTask)
        .set({
          status: "pending",
          startedAt: null,
          completedAt: null,
          updatedAt: new Date(),
        })
        .where(
          and(
            eq(canvasImageTask.id, task.id),
            eq(canvasImageTask.status, "running"),
            lt(canvasImageTask.startedAt, staleBefore)
        )
        );
      await enqueueLocalTask(task.id, {
        modelKey: normalizeTaskModelKey(task.request),
        userId: params.userId,
      });
      await drainCanvasImageTaskQueue();
      task.status = "pending" as CanvasImageTaskStatus;
      task.startedAt = null;
      task.completedAt = null;
    }
  }

  if (task.status === "pending") {
    await enqueueLocalTask(task.id, {
      modelKey: normalizeTaskModelKey(task.request),
      userId: params.userId,
    });
    await drainCanvasImageTaskQueue();
  }

  return task;
}