74 lines
2.0 KiB
TypeScript
74 lines
2.0 KiB
TypeScript
type GenerationKind = "image" | "video";
|
|
|
|
interface GenerationSlot {
|
|
id: string;
|
|
userKey: string;
|
|
kind: GenerationKind;
|
|
createdAt: number;
|
|
}
|
|
|
|
const DEFAULT_MAX_ACTIVE_GENERATION_TASKS = 3;
|
|
const STALE_SLOT_MS = 6 * 60 * 60 * 1000;
|
|
const activeSlots = new Map<string, GenerationSlot>();
|
|
|
|
let userMaxConcurrency: number | null = null;
|
|
|
|
export function setUserMaxConcurrency(limit: number | null | undefined): void {
|
|
userMaxConcurrency = typeof limit === "number" && limit > 0 ? limit : null;
|
|
}
|
|
|
|
function getEffectiveLimit(): number {
|
|
return userMaxConcurrency ?? DEFAULT_MAX_ACTIVE_GENERATION_TASKS;
|
|
}
|
|
|
|
export function getGenerationUserKey(userId?: string | number | null): string {
|
|
return userId === undefined || userId === null || userId === "" ? "anonymous" : String(userId);
|
|
}
|
|
|
|
function pruneStaleSlots(now = Date.now()): void {
|
|
activeSlots.forEach((slot, id) => {
|
|
if (now - slot.createdAt > STALE_SLOT_MS) {
|
|
activeSlots.delete(id);
|
|
}
|
|
});
|
|
}
|
|
|
|
export function getActiveGenerationTaskCount(userKey: string): number {
|
|
pruneStaleSlots();
|
|
let count = 0;
|
|
activeSlots.forEach((slot) => {
|
|
if (slot.userKey === userKey) count += 1;
|
|
});
|
|
return count;
|
|
}
|
|
|
|
export function claimGenerationSlot(input: {
|
|
userKey: string;
|
|
kind: GenerationKind;
|
|
id?: string;
|
|
}): () => void {
|
|
pruneStaleSlots();
|
|
const activeCount = getActiveGenerationTaskCount(input.userKey);
|
|
const effectiveLimit = getEffectiveLimit();
|
|
if (activeCount >= effectiveLimit) {
|
|
throw new Error(`当前账号同时最多生成 ${effectiveLimit} 个图片/视频任务，请等待已有任务完成后再提交。`);
|
|
}
|
|
|
|
const id = input.id || `generation-slot-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
|
activeSlots.set(id, {
|
|
id,
|
|
userKey: input.userKey,
|
|
kind: input.kind,
|
|
createdAt: Date.now(),
|
|
});
|
|
|
|
return () => {
|
|
activeSlots.delete(id);
|
|
};
|
|
}
|
|
|
|
export function releaseGenerationSlot(id: string | undefined | null): void {
|
|
if (!id) return;
|
|
activeSlots.delete(id);
|
|
}
|