931 lines
27 KiB
TypeScript
931 lines
27 KiB
TypeScript
import "server-only";
|
||
|
||
import { and, asc, desc, eq, lt, notInArray, sql } from "drizzle-orm";
|
||
import { db } from "@/lib/db";
|
||
import { canvasImageTask } from "@/lib/db/schema";
|
||
import {
|
||
getCanvasImageBatchDispatchIntervalMs,
|
||
getCanvasImageRunningCountForModel,
|
||
getCanvasImageSingleModelConcurrency,
|
||
normalizeCanvasImageModelKey,
|
||
} from "./image-concurrency";
|
||
import {
|
||
generateCanvasImage,
|
||
type CanvasGenerateImageRequest,
|
||
type CanvasGenerateImageResult,
|
||
} from "./generate-image";
|
||
import { applyScaleAdjustModelRelight } from "./scale-adjust-model-relight";
|
||
import {
|
||
executeStyleProTask,
|
||
getStyleProTaskPrimaryModelKey,
|
||
type StyleProTaskRequest,
|
||
type StyleProTaskResult,
|
||
} from "@/lib/style-pro/task";
|
||
import {
|
||
executeCanvasLayoutReplaceTask,
|
||
isCanvasLayoutReplaceTaskRequest,
|
||
type CanvasLayoutReplaceTaskRequest,
|
||
} from "@/lib/canvas/layout-replace-task";
|
||
import {
|
||
completeCanvasLayoutReplaceBatchItem,
|
||
markCanvasLayoutReplaceBatchItemRunning,
|
||
} from "@/lib/canvas/layout-replace-batch";
|
||
import { persistAgentProjectGeneratedResult } from "@/lib/agents/agent-project-result-persistence";
|
||
import { getUserFacingImageGenerationErrorMessage } from "@/lib/request-error";
|
||
|
||
type CanvasImageTaskStatus =
|
||
| "waiting"
|
||
| "pending"
|
||
| "running"
|
||
| "succeeded"
|
||
| "failed";
|
||
type CanvasImageTaskRequest =
|
||
| CanvasGenerateImageRequest
|
||
| StyleProTaskRequest
|
||
| CanvasLayoutReplaceTaskRequest;
|
||
type CanvasImageTaskResult = CanvasGenerateImageResult | StyleProTaskResult;
|
||
|
||
export type CanvasImageTaskQueueStatus = {
|
||
status: CanvasImageTaskStatus;
|
||
queuePosition: number;
|
||
sameModelPendingCount: number;
|
||
sameModelRunningCount: number;
|
||
concurrencyLimit: number;
|
||
estimatedWaitMs: number;
|
||
source: "model_queue" | "batch_queue";
|
||
};
|
||
|
||
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>;
|
||
lastStaleRunningCleanupAt: number;
|
||
draining: boolean;
|
||
}
|
||
|
||
const DEFAULT_RUNNING_TIMEOUT_MS = 10 * 60 * 1000;
|
||
const SCALE_ADJUST_RUNNING_TIMEOUT_MS = 12 * 60 * 1000;
|
||
const LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS = 30 * 60 * 1000;
|
||
const STALE_RUNNING_CLEANUP_INTERVAL_MS = 60 * 1000;
|
||
const PENDING_LOOKAHEAD_LIMIT = 100;
|
||
const QUEUE_STATUS_PENDING_LOOKAHEAD_LIMIT = 500;
|
||
const QUEUE_STATUS_RECENT_SAMPLE_LIMIT = 120;
|
||
const DEFAULT_ESTIMATED_WAIT_MS = 45 * 1000;
|
||
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>(),
|
||
lastStaleRunningCleanupAt: 0,
|
||
draining: false,
|
||
};
|
||
}
|
||
return globalObj[queueStateKey] as QueueState;
|
||
}
|
||
|
||
function normalizeTaskModelKey(request: unknown) {
|
||
if (isStyleProTaskRequest(request)) {
|
||
return normalizeCanvasImageModelKey(getStyleProTaskPrimaryModelKey(request));
|
||
}
|
||
|
||
if (!request || typeof request !== "object" || Array.isArray(request)) {
|
||
return "__default__";
|
||
}
|
||
|
||
const rawModel = (request as Record<string, unknown>).model;
|
||
return normalizeCanvasImageModelKey(rawModel);
|
||
}
|
||
|
||
function getRecord(value: unknown): Record<string, unknown> | null {
|
||
return value && typeof value === "object" && !Array.isArray(value)
|
||
? (value as Record<string, unknown>)
|
||
: null;
|
||
}
|
||
|
||
function getCanvasImageRunningTimeoutMs(request: unknown) {
|
||
const record = getRecord(request);
|
||
const debugMeta = getRecord(record?.debugMeta);
|
||
const debugLabel = String(record?.debugLabel || "");
|
||
const operation = String(debugMeta?.operation || "");
|
||
const isScaleAdjust =
|
||
debugLabel === "canvas-scale-adjust" ||
|
||
operation === "scale-adjust" ||
|
||
Boolean(debugMeta?.scaleAdjustPatchWorkflowApplied) ||
|
||
Boolean(debugMeta?.scaleAdjustPatchForcedForShrink) ||
|
||
Boolean(record?.scaleAdjustPatchComposite);
|
||
|
||
return isScaleAdjust
|
||
? SCALE_ADJUST_RUNNING_TIMEOUT_MS
|
||
: DEFAULT_RUNNING_TIMEOUT_MS;
|
||
}
|
||
|
||
async function withCanvasImageTaskExecutionTimeout<T>(
|
||
operation: Promise<T>,
|
||
timeoutMs: number,
|
||
taskId: string
|
||
): Promise<T> {
|
||
if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
|
||
return operation;
|
||
}
|
||
|
||
let timeoutId: ReturnType<typeof setTimeout> | null = null;
|
||
try {
|
||
return await Promise.race([
|
||
operation,
|
||
new Promise<T>((_, reject) => {
|
||
timeoutId = setTimeout(() => {
|
||
reject(
|
||
new Error(
|
||
`图片生成任务执行超时（${Math.round(timeoutMs / 1000)}秒，${taskId}）`
|
||
)
|
||
);
|
||
}, timeoutMs);
|
||
}),
|
||
]);
|
||
} finally {
|
||
if (timeoutId) {
|
||
clearTimeout(timeoutId);
|
||
}
|
||
}
|
||
}
|
||
|
||
function isStyleProTaskRequest(request: unknown): request is StyleProTaskRequest {
|
||
return Boolean(
|
||
request &&
|
||
typeof request === "object" &&
|
||
!Array.isArray(request) &&
|
||
(request as { taskKind?: unknown }).taskKind === "style-pro"
|
||
);
|
||
}
|
||
|
||
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 markStaleDatabaseRunningTasksFailed() {
|
||
const state = getQueueState();
|
||
const now = Date.now();
|
||
if (now - state.lastStaleRunningCleanupAt < STALE_RUNNING_CLEANUP_INTERVAL_MS) {
|
||
return;
|
||
}
|
||
state.lastStaleRunningCleanupAt = now;
|
||
|
||
const localRunningIds = Array.from(state.running.keys());
|
||
const staleBefore = new Date(now - LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS);
|
||
const staleWhere =
|
||
localRunningIds.length > 0
|
||
? and(
|
||
eq(canvasImageTask.status, "running"),
|
||
lt(canvasImageTask.startedAt, staleBefore),
|
||
notInArray(canvasImageTask.id, localRunningIds)
|
||
)
|
||
: and(
|
||
eq(canvasImageTask.status, "running"),
|
||
lt(canvasImageTask.startedAt, staleBefore)
|
||
);
|
||
|
||
const completedAt = new Date(now);
|
||
const staleTasks = await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "failed",
|
||
error: "图片生成任务已超时，请重新发起生成",
|
||
completedAt,
|
||
updatedAt: completedAt,
|
||
})
|
||
.where(staleWhere)
|
||
.returning({ id: canvasImageTask.id });
|
||
|
||
if (staleTasks.length > 0) {
|
||
console.warn("[Canvas Image Task Queue] Marked stale running tasks failed:", {
|
||
count: staleTasks.length,
|
||
taskIds: staleTasks.slice(0, 10).map((task) => task.id),
|
||
});
|
||
}
|
||
}
|
||
|
||
async function runSingleTask(entry: PendingTaskEntry) {
|
||
let completedBatchRequest: CanvasLayoutReplaceTaskRequest | null = null;
|
||
let completedBatchResult: CanvasImageTaskResult | null = null;
|
||
let completedBatchError: string | null = null;
|
||
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 request = task.request as CanvasImageTaskRequest;
|
||
if (isCanvasLayoutReplaceTaskRequest(request) && request.batchId) {
|
||
completedBatchRequest = request;
|
||
try {
|
||
await markCanvasLayoutReplaceBatchItemRunning(request);
|
||
} catch (error) {
|
||
console.error("[Canvas Image Task Queue] Failed to mark batch item running:", {
|
||
batchId: request.batchId,
|
||
error: error instanceof Error ? error.message : String(error),
|
||
taskId: entry.taskId,
|
||
});
|
||
}
|
||
}
|
||
const runningTimeoutMs = getCanvasImageRunningTimeoutMs(request);
|
||
const execution: Promise<CanvasImageTaskResult> = isStyleProTaskRequest(request)
|
||
? executeStyleProTask({
|
||
request,
|
||
userId: String(task.userId || entry.userId),
|
||
})
|
||
: isCanvasLayoutReplaceTaskRequest(request)
|
||
? executeCanvasLayoutReplaceTask({
|
||
request,
|
||
userId: String(task.userId || entry.userId),
|
||
})
|
||
: (async () => {
|
||
const imageRequest = request as CanvasGenerateImageRequest;
|
||
const ownerKey = String(task.userId || entry.userId);
|
||
const initialResult = await generateCanvasImage({
|
||
...imageRequest,
|
||
ownerKey,
|
||
});
|
||
return await applyScaleAdjustModelRelight({
|
||
generateImage: generateCanvasImage,
|
||
initialResult,
|
||
ownerKey,
|
||
request: imageRequest,
|
||
});
|
||
})();
|
||
const result = await withCanvasImageTaskExecutionTimeout(
|
||
execution,
|
||
runningTimeoutMs,
|
||
entry.taskId
|
||
);
|
||
completedBatchResult = result;
|
||
if (isCanvasLayoutReplaceTaskRequest(request) && request.agentProjectId) {
|
||
const layoutReplaceResult = result as CanvasGenerateImageResult;
|
||
const generatedImageUrl = String(
|
||
layoutReplaceResult.images?.[0]?.originalUrl ||
|
||
layoutReplaceResult.images?.[0]?.thumbnailUrl ||
|
||
""
|
||
).trim();
|
||
try {
|
||
await persistAgentProjectGeneratedResult({
|
||
projectId: request.agentProjectId,
|
||
userId: String(task.userId || entry.userId),
|
||
imageUrl: generatedImageUrl,
|
||
modeType: "scene_generation",
|
||
});
|
||
} catch (error) {
|
||
console.error("[Canvas Image Task Queue] Failed to persist agent AI template result:", {
|
||
error: error instanceof Error ? error.message : String(error),
|
||
projectId: request.agentProjectId,
|
||
taskId: entry.taskId,
|
||
});
|
||
}
|
||
}
|
||
await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "succeeded",
|
||
result: result as CanvasImageTaskResult,
|
||
error: null,
|
||
completedAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(
|
||
and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "running"))
|
||
);
|
||
} catch (error) {
|
||
console.error("[Canvas Image Task Queue] Image task failed:", {
|
||
taskId: entry.taskId,
|
||
error,
|
||
});
|
||
completedBatchError = getUserFacingImageGenerationErrorMessage(error);
|
||
await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "failed",
|
||
error: completedBatchError,
|
||
completedAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(
|
||
and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "running"))
|
||
);
|
||
} finally {
|
||
const modelKey = releaseRunningTask(entry.taskId);
|
||
let nextBatchTaskId: string | null = null;
|
||
if (completedBatchRequest?.batchId) {
|
||
try {
|
||
nextBatchTaskId = await completeCanvasLayoutReplaceBatchItem({
|
||
request: completedBatchRequest,
|
||
result: completedBatchResult as CanvasGenerateImageResult | null,
|
||
error: completedBatchError,
|
||
});
|
||
} catch (error) {
|
||
console.error("[Canvas Image Task Queue] Failed to advance layout batch:", {
|
||
batchId: completedBatchRequest.batchId,
|
||
error: error instanceof Error ? error.message : String(error),
|
||
taskId: entry.taskId,
|
||
});
|
||
}
|
||
}
|
||
if (modelKey) {
|
||
const batchDispatchIntervalMs = await getCanvasImageBatchDispatchIntervalMs();
|
||
if (getQueueState().awaitingNextBatch.has(modelKey)) {
|
||
scheduleModelNextBatch(modelKey, batchDispatchIntervalMs);
|
||
}
|
||
}
|
||
if (nextBatchTaskId) {
|
||
await enqueueLocalTask(nextBatchTaskId);
|
||
}
|
||
void drainCanvasImageTaskQueue();
|
||
}
|
||
}
|
||
|
||
export async function drainCanvasImageTaskQueue() {
|
||
const state = getQueueState();
|
||
if (state.draining) return;
|
||
|
||
state.draining = true;
|
||
try {
|
||
await markStaleDatabaseRunningTasksFailed();
|
||
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 ensureCanvasImageTaskActive(params: {
|
||
taskId: string;
|
||
userId: string;
|
||
}) {
|
||
const [task] = await db
|
||
.select({
|
||
id: canvasImageTask.id,
|
||
status: canvasImageTask.status,
|
||
request: canvasImageTask.request,
|
||
result: canvasImageTask.result,
|
||
error: canvasImageTask.error,
|
||
})
|
||
.from(canvasImageTask)
|
||
.where(
|
||
and(
|
||
eq(canvasImageTask.id, params.taskId),
|
||
eq(canvasImageTask.userId, params.userId)
|
||
)
|
||
)
|
||
.limit(1);
|
||
if (!task || task.status === "waiting") return false;
|
||
|
||
const request = task.request as CanvasImageTaskRequest;
|
||
if (
|
||
(task.status === "succeeded" || task.status === "failed") &&
|
||
isCanvasLayoutReplaceTaskRequest(request) &&
|
||
request.batchId
|
||
) {
|
||
const nextTaskId = await completeCanvasLayoutReplaceBatchItem({
|
||
request,
|
||
result: task.result as CanvasGenerateImageResult | null,
|
||
error: task.status === "failed" ? task.error || "排版替换失败" : null,
|
||
});
|
||
if (nextTaskId) {
|
||
await enqueueLocalTask(nextTaskId);
|
||
await drainCanvasImageTaskQueue();
|
||
}
|
||
return true;
|
||
}
|
||
|
||
if (task.status === "running" && !getQueueState().running.has(task.id)) {
|
||
await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "pending",
|
||
startedAt: null,
|
||
completedAt: null,
|
||
error: null,
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(
|
||
and(
|
||
eq(canvasImageTask.id, task.id),
|
||
eq(canvasImageTask.status, "running")
|
||
)
|
||
);
|
||
} else if (task.status !== "pending" && task.status !== "running") {
|
||
return false;
|
||
}
|
||
|
||
await enqueueLocalTask(task.id, { userId: params.userId });
|
||
await drainCanvasImageTaskQueue();
|
||
return true;
|
||
}
|
||
|
||
export async function createCanvasImageTask(params: {
|
||
userId: string;
|
||
request: Record<string, unknown>;
|
||
}) {
|
||
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;
|
||
}
|
||
|
||
async function estimateModelQueueWaitMs(params: {
|
||
modelKey: string;
|
||
sameModelRunningCount: number;
|
||
queuePosition: number;
|
||
concurrencyLimit: number;
|
||
}) {
|
||
const rows = await db
|
||
.select({
|
||
request: canvasImageTask.request,
|
||
startedAt: canvasImageTask.startedAt,
|
||
completedAt: canvasImageTask.completedAt,
|
||
})
|
||
.from(canvasImageTask)
|
||
.where(eq(canvasImageTask.status, "succeeded"))
|
||
.orderBy(desc(canvasImageTask.completedAt))
|
||
.limit(QUEUE_STATUS_RECENT_SAMPLE_LIMIT);
|
||
|
||
const durations = rows
|
||
.filter((row) => normalizeTaskModelKey(row.request) === params.modelKey)
|
||
.map((row) => {
|
||
if (!row.startedAt || !row.completedAt) return 0;
|
||
return row.completedAt.getTime() - row.startedAt.getTime();
|
||
})
|
||
.filter((value) => Number.isFinite(value) && value > 1000)
|
||
.slice(0, 12);
|
||
|
||
const averageDurationMs =
|
||
durations.length > 0
|
||
? Math.round(durations.reduce((sum, value) => sum + value, 0) / durations.length)
|
||
: DEFAULT_ESTIMATED_WAIT_MS;
|
||
const immediatelyStartable = Math.max(
|
||
params.concurrencyLimit - params.sameModelRunningCount,
|
||
0
|
||
);
|
||
|
||
if (params.queuePosition < immediatelyStartable) {
|
||
return 0;
|
||
}
|
||
|
||
const tasksAheadAfterImmediateSlots = Math.max(
|
||
0,
|
||
params.queuePosition - immediatelyStartable
|
||
);
|
||
const waveCount =
|
||
Math.floor(
|
||
tasksAheadAfterImmediateSlots / Math.max(1, params.concurrencyLimit)
|
||
) +
|
||
(params.sameModelRunningCount > 0 || immediatelyStartable === 0 ? 1 : 0);
|
||
|
||
return Math.max(0, waveCount * averageDurationMs);
|
||
}
|
||
|
||
async function buildCanvasImageTaskQueueStatus(task: {
|
||
id: string;
|
||
status: CanvasImageTaskStatus;
|
||
request: unknown;
|
||
createdAt: Date;
|
||
}) {
|
||
const modelKey = normalizeTaskModelKey(task.request);
|
||
const concurrencyLimit = await getCanvasImageSingleModelConcurrency();
|
||
const runningRows = await db
|
||
.select({
|
||
request: canvasImageTask.request,
|
||
})
|
||
.from(canvasImageTask)
|
||
.where(eq(canvasImageTask.status, "running"));
|
||
const sameModelRunningCount = Math.max(
|
||
runningRows.filter((row) => normalizeTaskModelKey(row.request) === modelKey).length,
|
||
getCanvasImageRunningCountForModel(modelKey)
|
||
);
|
||
|
||
if (task.status !== "pending") {
|
||
return {
|
||
status: task.status,
|
||
queuePosition: 0,
|
||
sameModelPendingCount: 0,
|
||
sameModelRunningCount,
|
||
concurrencyLimit,
|
||
estimatedWaitMs: 0,
|
||
source: "model_queue" as const,
|
||
};
|
||
}
|
||
|
||
const pendingRows = await db
|
||
.select({
|
||
id: canvasImageTask.id,
|
||
request: canvasImageTask.request,
|
||
createdAt: canvasImageTask.createdAt,
|
||
})
|
||
.from(canvasImageTask)
|
||
.where(eq(canvasImageTask.status, "pending"))
|
||
.orderBy(asc(canvasImageTask.createdAt))
|
||
.limit(QUEUE_STATUS_PENDING_LOOKAHEAD_LIMIT);
|
||
|
||
const sameModelPendingRows = pendingRows.filter(
|
||
(row) => normalizeTaskModelKey(row.request) === modelKey
|
||
);
|
||
const queuePosition = Math.max(
|
||
sameModelPendingRows.findIndex((row) => row.id === task.id),
|
||
0
|
||
);
|
||
const sameModelPendingCount = sameModelPendingRows.length;
|
||
const estimatedWaitMs = await estimateModelQueueWaitMs({
|
||
modelKey,
|
||
sameModelRunningCount,
|
||
queuePosition,
|
||
concurrencyLimit,
|
||
});
|
||
|
||
return {
|
||
status: task.status,
|
||
queuePosition,
|
||
sameModelPendingCount,
|
||
sameModelRunningCount,
|
||
concurrencyLimit,
|
||
estimatedWaitMs,
|
||
source: "model_queue" as const,
|
||
};
|
||
}
|
||
|
||
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 now = Date.now();
|
||
const runningAgeMs = now - task.startedAt.getTime();
|
||
const runningTimeoutMs = getCanvasImageRunningTimeoutMs(task.request);
|
||
const locallyRunning = getQueueState().running.has(task.id);
|
||
const shouldMarkStale =
|
||
runningAgeMs >= runningTimeoutMs &&
|
||
(!locallyRunning || runningAgeMs >= LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS);
|
||
|
||
if (shouldMarkStale) {
|
||
const staleBefore = new Date(now - runningTimeoutMs);
|
||
const timeoutMessage = "图片生成任务已超时，请重新发起生成";
|
||
const completedAt = new Date();
|
||
const [timedOutTask] = await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "failed",
|
||
error: timeoutMessage,
|
||
completedAt,
|
||
updatedAt: completedAt,
|
||
})
|
||
.where(
|
||
and(
|
||
eq(canvasImageTask.id, task.id),
|
||
eq(canvasImageTask.status, "running"),
|
||
lt(canvasImageTask.startedAt, staleBefore)
|
||
)
|
||
)
|
||
.returning({
|
||
status: canvasImageTask.status,
|
||
error: canvasImageTask.error,
|
||
completedAt: canvasImageTask.completedAt,
|
||
});
|
||
|
||
if (timedOutTask) {
|
||
releaseRunningTask(task.id);
|
||
task.status = timedOutTask.status as CanvasImageTaskStatus;
|
||
task.error = timedOutTask.error;
|
||
task.completedAt = timedOutTask.completedAt;
|
||
}
|
||
}
|
||
}
|
||
|
||
if (task.status === "pending") {
|
||
await enqueueLocalTask(task.id, {
|
||
modelKey: normalizeTaskModelKey(task.request),
|
||
userId: params.userId,
|
||
});
|
||
await drainCanvasImageTaskQueue();
|
||
}
|
||
|
||
return {
|
||
...task,
|
||
queueStatus: await buildCanvasImageTaskQueueStatus({
|
||
id: task.id,
|
||
status: task.status as CanvasImageTaskStatus,
|
||
request: task.request,
|
||
createdAt: task.createdAt,
|
||
}),
|
||
};
|
||
}
|