All BatchedQueues should drain when disposed (also adds query string for output: "debug=true") (#5098)

* ensure disposal is correct

* add tests for disposal of batching queue

* fixes for draining disposal

* Update packages/objectloader2/src/queues/batchingQueue.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix comment

* fix tests and build

* add query string inspection of debug parameter

* Update packages/objectloader2/src/queues/batchingQueue.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update packages/objectloader2/src/core/objectLoader2Factory.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix test

* fix AI

* export getQueryParameter to avoid dup code.  Sandbox uses it too

* add tests for functions

* prettier fix

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

packages/objectloader2/src/queues/batchingQueue.dispose.test.ts

@@ -0,0 +1,74 @@
+import { describe, test, expect, vi } from 'vitest'
+import BatchingQueue from './batchingQueue.js'
+
+describe('BatchingQueue disposal', () => {
+  test('should drain the queue on dispose', async () => {
+    const processFunction = vi.fn().mockResolvedValue(undefined)
+    const queue = new BatchingQueue<{ id: string }>({
+      batchSize: 5,
+      maxWaitTime: 1000,
+      processFunction
+    })
+
+    const items = Array.from({ length: 3 }, (_, i) => ({ id: `item-${i}` }))
+    items.forEach((item) => queue.add(item.id, item))
+
+    expect(queue.count()).toBe(3)
+
+    await queue.disposeAsync()
+
+    expect(processFunction).toHaveBeenCalledWith(items)
+    expect(queue.count()).toBe(0)
+    expect(queue.isDisposed()).toBe(true)
+  })
+
+  test('should wait for processing to finish before disposing', async () => {
+    let resolveProcess: (value: void | PromiseLike<void>) => void = () => {}
+    const processPromise = new Promise<void>((resolve) => {
+      resolveProcess = resolve
+    })
+
+    const processFunction = vi.fn().mockImplementation(() => processPromise)
+
+    const queue = new BatchingQueue<{ id: string }>({
+      batchSize: 2,
+      maxWaitTime: 100,
+      processFunction
+    })
+
+    const items1 = [{ id: 'item-1' }, { id: 'item-2' }]
+    items1.forEach((item) => queue.add(item.id, item))
+
+    // First batch is processing
+    expect(processFunction).toHaveBeenCalledWith(items1)
+
+    const items2 = [{ id: 'item-3' }]
+    items2.forEach((item) => queue.add(item.id, item))
+
+    const disposePromise = queue.disposeAsync()
+
+    // Queue should be disposed now, but processing is still ongoing
+    expect(queue.isDisposed()).toBe(true)
+    resolveProcess()
+    await disposePromise
+
+    expect(processFunction).toHaveBeenCalledTimes(2)
+    expect(processFunction).toHaveBeenCalledWith(items2)
+    expect(queue.count()).toBe(0)
+    expect(queue.isDisposed()).toBe(true)
+  })
+
+  test('adding items after dispose should do nothing', async () => {
+    const processFunction = vi.fn().mockResolvedValue(undefined)
+    const queue = new BatchingQueue<string>({
+      batchSize: 5,
+      maxWaitTime: 1000,
+      processFunction
+    })
+
+    await queue.disposeAsync()
+    queue.add('key1', 'item1')
+    expect(queue.count()).toBe(0)
+    expect(processFunction).not.toHaveBeenCalled()
+  })
+})

packages/objectloader2/src/queues/batchingQueue.test.ts

@@ -1,26 +1,10 @@
-import { describe, test, expect, beforeEach, afterEach, vi } from 'vitest'
+import { describe, test, expect, vi } from 'vitest'
 import BatchingQueue from './batchingQueue.js'
 
 describe('BatchingQueue', () => {
-  let queue: BatchingQueue<string>
-
-  beforeEach(() => {
-    queue = new BatchingQueue({
-      batchSize: 3,
-      maxWaitTime: 100,
-      processFunction: async (): Promise<void> => {
-        await new Promise((resolve) => setTimeout(resolve, 0))
-      }
-    })
-  })
-
-  afterEach(() => {
-    queue.dispose()
-  })
-
   test('should add items and process them in batches', async () => {
     const processSpy = vi.fn()
-    queue = new BatchingQueue({
+    const queue = new BatchingQueue({
       batchSize: 2,
       maxWaitTime: 100,
       processFunction: async (batch: string[]): Promise<void> => {
@@ -29,18 +13,22 @@ describe('BatchingQueue', () => {
       }
     })
 
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
+    try {
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
 
-    await new Promise((resolve) => setTimeout(resolve, 200))
+      await new Promise((resolve) => setTimeout(resolve, 200))
 
-    expect(processSpy).toHaveBeenCalledTimes(1)
-    expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
+      expect(processSpy).toHaveBeenCalledTimes(1)
+      expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 
   test('should process items after timeout if batch size is not reached', async () => {
     const processSpy = vi.fn()
-    queue = new BatchingQueue({
+    const queue = new BatchingQueue({
       batchSize: 5,
       maxWaitTime: 100,
       processFunction: async (batch: string[]): Promise<void> => {
@@ -49,37 +37,22 @@ describe('BatchingQueue', () => {
       }
     })
 
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
+    try {
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
 
-    await new Promise((resolve) => setTimeout(resolve, 200))
+      await new Promise((resolve) => setTimeout(resolve, 200))
 
-    expect(processSpy).toHaveBeenCalledTimes(1)
-    expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
-  })
-
-  test('should not process items if disposed', async () => {
-    const processSpy = vi.fn()
-    queue = new BatchingQueue({
-      batchSize: 2,
-      maxWaitTime: 10000,
-      processFunction: async (batch: string[]): Promise<void> => {
-        await new Promise((resolve) => setTimeout(resolve, 0))
-        processSpy(batch)
-      }
-    })
-
-    queue.add('key1', 'item1')
-    queue.dispose()
-
-    await new Promise((resolve) => setTimeout(resolve, 200))
-
-    expect(processSpy).not.toHaveBeenCalled()
+      expect(processSpy).toHaveBeenCalledTimes(1)
+      expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 
   test('should handle multiple batches correctly', async () => {
     const processSpy = vi.fn()
-    queue = new BatchingQueue({
+    const queue = new BatchingQueue({
       batchSize: 2,
       maxWaitTime: 100,
       processFunction: async (batch: string[]): Promise<void> => {
@@ -88,39 +61,65 @@ describe('BatchingQueue', () => {
       }
     })
 
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
-    queue.add('key3', 'item3')
-    queue.add('key4', 'item4')
+    try {
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
+      queue.add('key3', 'item3')
+      queue.add('key4', 'item4')
 
-    await new Promise((resolve) => setTimeout(resolve, 200))
+      await new Promise((resolve) => setTimeout(resolve, 200))
 
-    expect(processSpy).toHaveBeenCalledTimes(2)
-    expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
-    expect(processSpy).toHaveBeenCalledWith(['item3', 'item4'])
+      expect(processSpy).toHaveBeenCalledTimes(2)
+      expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
+      expect(processSpy).toHaveBeenCalledWith(['item3', 'item4'])
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 
-  test('should retrieve items by key', () => {
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
+  test('should retrieve items by key', async () => {
+    const queue = new BatchingQueue<string>({
+      batchSize: 3,
+      maxWaitTime: 100,
+      processFunction: async (): Promise<void> => {
+        await new Promise((resolve) => setTimeout(resolve, 0))
+      }
+    })
+    try {
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
 
-    expect(queue.get('key1')).toBe('item1')
-    expect(queue.get('key2')).toBe('item2')
-    expect(queue.get('key3')).toBeUndefined()
+      expect(queue.get('key1')).toBe('item1')
+      expect(queue.get('key2')).toBe('item2')
+      expect(queue.get('key3')).toBeUndefined()
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 
-  test('should return correct count of items', () => {
-    expect(queue.count()).toBe(0)
+  test('should return correct count of items', async () => {
+    const queue = new BatchingQueue<string>({
+      batchSize: 3,
+      maxWaitTime: 100,
+      processFunction: async (): Promise<void> => {
+        await new Promise((resolve) => setTimeout(resolve, 0))
+      }
+    })
+    try {
+      expect(queue.count()).toBe(0)
 
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
 
-    expect(queue.count()).toBe(2)
+      expect(queue.count()).toBe(2)
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 
   test('should not process items if already processing', async () => {
     const processSpy = vi.fn()
-    queue = new BatchingQueue({
+    const queue = new BatchingQueue({
       batchSize: 2,
       maxWaitTime: 100,
       processFunction: async (batch: string[]): Promise<void> => {
@@ -129,18 +128,22 @@ describe('BatchingQueue', () => {
       }
     })
 
-    queue.add('key1', 'item1')
-    queue.add('key2', 'item2')
-    queue.add('key3', 'item3')
+    try {
+      queue.add('key1', 'item1')
+      queue.add('key2', 'item2')
+      queue.add('key3', 'item3')
 
-    await new Promise((resolve) => setTimeout(resolve, 200))
+      await new Promise((resolve) => setTimeout(resolve, 200))
 
-    expect(processSpy).toHaveBeenCalledTimes(1)
-    expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
+      expect(processSpy).toHaveBeenCalledTimes(1)
+      expect(processSpy).toHaveBeenCalledWith(['item1', 'item2'])
 
-    await new Promise((resolve) => setTimeout(resolve, 200))
+      await new Promise((resolve) => setTimeout(resolve, 200))
 
-    expect(processSpy).toHaveBeenCalledTimes(2)
-    expect(processSpy).toHaveBeenCalledWith(['item3'])
+      expect(processSpy).toHaveBeenCalledTimes(2)
+      expect(processSpy).toHaveBeenCalledWith(['item3'])
+    } finally {
+      await queue.disposeAsync()
+    }
   })
 })

packages/objectloader2/src/queues/batchingQueue.ts

@@ -1,6 +1,12 @@
 import { CustomLogger } from '../types/functions.js'
 import KeyedQueue from './keyedQueue.js'
 
+/**
+ * Default wait time in milliseconds for processing ongoing tasks during disposal.
+ * This value was chosen to balance responsiveness and CPU usage in typical scenarios.
+ */
+const PROCESSING_WAIT_TIME_MS = 100
+
 export default class BatchingQueue<T> {
   #queue: KeyedQueue<string, T> = new KeyedQueue<string, T>()
   #batchSize: number
@@ -43,24 +49,41 @@ export default class BatchingQueue<T> {
     this.#logger = params.logger || ((): void => {})
   }
 
-  dispose(): void {
+  async disposeAsync(): Promise<void> {
     this.#disposed = true
     if (this.#timeoutId) {
       this.#getClearTimeoutFn()(this.#timeoutId)
+      this.#timeoutId = null
+    }
+
+    // Wait for any ongoing processing to finish
+    while (this.#isProcessing) {
+      await new Promise((resolve) =>
+        this.#getSetTimeoutFn()(resolve, PROCESSING_WAIT_TIME_MS)
+      )
+    }
+
+    // After any ongoing flush is completed, there might be items in the queue.
+    // We should flush them.
+    if (this.#queue.size > 0) {
+      await this.#flush()
     }
   }
 
   add(key: string, item: T): void {
+    if (this.#disposed) return
     this.#queue.enqueue(key, item)
     this.#addCheck()
   }
 
   addAll(keys: string[], items: T[]): void {
+    if (this.#disposed) return
     this.#queue.enqueueAll(keys, items)
     this.#addCheck()
   }
 
   #addCheck(): void {
+    if (this.#disposed) return
     if (this.#queue.size >= this.#batchSize) {
       // Fire and forget, no need to await
       // eslint-disable-next-line @typescript-eslint/no-floating-promises