63 lines
1.6 KiB
TypeScript
63 lines
1.6 KiB
TypeScript
|
type GenerationKind = "image" | "video";
|
||
|
|
|
||
|
|
interface GenerationSlot {
|
||
|
|
id: string;
|
||
|
|
userKey: string;
|
||
|
|
kind: GenerationKind;
|
||
|
|
createdAt: number;
|
||
|
|
}
|
||
|
|
|
||
|
|
const MAX_ACTIVE_GENERATION_TASKS = 3;
|
||
|
|
const STALE_SLOT_MS = 6 * 60 * 60 * 1000;
|
||
|
|
const activeSlots = new Map<string, GenerationSlot>();
|
||
|
|
|
||
|
|
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);
|
||
|
|
if (activeCount >= MAX_ACTIVE_GENERATION_TASKS) {
|
||
|
|
throw new Error("当前账号同时最多生成 3 个图片/视频任务，请等待已有任务完成后再提交。");
|
||
|
|
}
|
||
|
|
|
||
|
|
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);
|
||
|
|
}
|