@dickpy/dsh-imagegen 1.3.0 → 1.5.0

package/src/task-queue.ts

@@ -9,9 +9,13 @@ export class GenerationTaskQueue {
   private readonly tasks: GenerationTask[] = []
   private readonly controllers = new Map<string, AbortController>()
   private readonly listeners = new Set<GenerationTaskListener>()
-  private running = false
+  private running = 0
+  private serialRunning = false
 
-  constructor(private readonly run: (request: GenerateRequest, signal: AbortSignal) => Promise<GenerateResult>) {}
+  constructor(
+    private readonly run: (request: GenerateRequest, signal: AbortSignal) => Promise<GenerateResult>,
+    private readonly concurrency = 1,
+  ) {}
 
   list(): GenerationTask[] {
     return this.tasks.map(task => this.snapshot(task))
@@ -27,7 +31,7 @@ export class GenerationTaskQueue {
     const task: GenerationTask = { id: randomUUID(), request: { ...request }, status: 'queued', createdAt: Date.now() }
     this.tasks.unshift(task)
     this.publish(task)
-    void this.drain()
+    this.drain()
     return this.snapshot(task)
   }
 
@@ -38,6 +42,7 @@ export class GenerationTaskQueue {
     task.finishedAt = Date.now()
     this.controllers.get(id)?.abort()
     this.publish(task)
+    this.drain()
     return this.snapshot(task)
   }
 
@@ -46,39 +51,44 @@ export class GenerationTaskQueue {
     return previous === undefined ? undefined : this.submit(previous.request)
   }
 
-  private async drain(): Promise<void> {
-    if (this.running) return
-    this.running = true
+  private drain(): void {
+    while (this.running < Math.max(1, this.concurrency)) {
+      const task = this.tasks.find(item => item.status === 'queued'
+        && (this.running === 0 || (item.request.comparisonId !== undefined && !this.serialRunning)))
+      if (task === undefined) return
+      this.running += 1
+      if (task.request.comparisonId === undefined) this.serialRunning = true
+      void this.runTask(task).finally(() => {
+        this.running -= 1
+        if (task.request.comparisonId === undefined) this.serialRunning = false
+        this.drain()
+      })
+    }
+  }
+
+  private async runTask(task: GenerationTask): Promise<void> {
+    task.status = 'running'
+    task.startedAt = Date.now()
+    this.publish(task)
+    const controller = new AbortController()
+    this.controllers.set(task.id, controller)
     try {
-      for (;;) {
-        const task = this.tasks.find(item => item.status === 'queued')
-        if (task === undefined) return
-        task.status = 'running'
-        task.startedAt = Date.now()
+      const result = await this.run(task.request, controller.signal)
+      if (this.tasks.find(item => item.id === task.id)?.status !== 'cancelled') {
+        task.status = 'completed'
+        task.result = result
+        task.finishedAt = Date.now()
+        this.publish(task)
+      }
+    } catch (error) {
+      if (this.tasks.find(item => item.id === task.id)?.status !== 'cancelled') {
+        task.status = 'failed'
+        task.error = error instanceof Error ? error.message : String(error)
+        task.finishedAt = Date.now()
         this.publish(task)
-        const controller = new AbortController()
-        this.controllers.set(task.id, controller)
-        try {
-          const result = await this.run(task.request, controller.signal)
-          if (this.tasks.find(item => item.id === task.id)?.status !== 'cancelled') {
-            task.status = 'completed'
-            task.result = result
-            task.finishedAt = Date.now()
-            this.publish(task)
-          }
-        } catch (error) {
-          if (this.tasks.find(item => item.id === task.id)?.status !== 'cancelled') {
-            task.status = 'failed'
-            task.error = error instanceof Error ? error.message : String(error)
-            task.finishedAt = Date.now()
-            this.publish(task)
-          }
-        } finally {
-          this.controllers.delete(task.id)
-        }
       }
     } finally {
-      this.running = false
+      this.controllers.delete(task.id)
     }
   }