ref(core): Improve promise buffer

Author: mydea

packages/core/src/utils/promisebuffer.ts

@@ -1,9 +1,9 @@
-import { rejectedSyncPromise, resolvedSyncPromise, SyncPromise } from './syncpromise';
+import { rejectedSyncPromise, resolvedSyncPromise } from './syncpromise';
 
 export interface PromiseBuffer<T> {
   // exposes the internal array so tests can assert on the state of it.
   // XXX: this really should not be public api.
-  $: Array<PromiseLike<T>>;
+  $: PromiseLike<T>[];
   add(taskProducer: () => PromiseLike<T>): PromiseLike<T>;
   drain(timeout?: number): PromiseLike<boolean>;
 }
@@ -14,11 +14,11 @@ export const SENTRY_BUFFER_FULL_ERROR = Symbol.for('SentryBufferFullError');
  * Creates an new PromiseBuffer object with the specified limit
  * @param limit max number of promises that can be stored in the buffer
  */
-export function makePromiseBuffer<T>(limit?: number): PromiseBuffer<T> {
-  const buffer: Array<PromiseLike<T>> = [];
+export function makePromiseBuffer<T>(limit: number = 100): PromiseBuffer<T> {
+  const buffer: Set<PromiseLike<T>> = new Set();
 
   function isReady(): boolean {
-    return limit === undefined || buffer.length < limit;
+    return buffer.size < limit;
   }
 
   /**
@@ -27,8 +27,8 @@ export function makePromiseBuffer<T>(limit?: number): PromiseBuffer<T> {
    * @param task Can be any PromiseLike<T>
    * @returns Removed promise.
    */
-  function remove(task: PromiseLike<T>): PromiseLike<T | void> {
-    return buffer.splice(buffer.indexOf(task), 1)[0] || Promise.resolve(undefined);
+  function remove(task: PromiseLike<T>): void {
+    buffer.delete(task);
   }
 
   /**
@@ -48,22 +48,26 @@ export function makePromiseBuffer<T>(limit?: number): PromiseBuffer<T> {
 
     // start the task and add its promise to the queue
     const task = taskProducer();
-    if (buffer.indexOf(task) === -1) {
-      buffer.push(task);
-    }
-    void task
-      .then(() => remove(task))
-      // Use `then(null, rejectionHandler)` rather than `catch(rejectionHandler)` so that we can use `PromiseLike`
-      // rather than `Promise`. `PromiseLike` doesn't have a `.catch` method, making its polyfill smaller. (ES5 didn't
-      // have promises, so TS has to polyfill when down-compiling.)
-      .then(null, () =>
-        remove(task).then(null, () => {
-          // We have to add another catch here because `remove()` starts a new promise chain.
-        }),
-      );
+    buffer.add(task);
+    void task.then(
+      () => remove(task),
+      () => remove(task),
+    );
     return task;
   }
 
+  function drainNextSyncPromise(): PromiseLike<boolean> {
+    const item = buffer.values().next().value;
+
+    if (!item) {
+      return resolvedSyncPromise(true);
+    }
+
+    return resolvedSyncPromise(item).then(() => {
+      return drainNextSyncPromise();
+    });
+  }
+
   /**
    * Wait for all promises in the queue to resolve or for timeout to expire, whichever comes first.
    *
@@ -74,34 +78,27 @@ export function makePromiseBuffer<T>(limit?: number): PromiseBuffer<T> {
    * `false` otherwise
    */
   function drain(timeout?: number): PromiseLike<boolean> {
-    return new SyncPromise<boolean>((resolve, reject) => {
-      let counter = buffer.length;
-
-      if (!counter) {
-        return resolve(true);
-      }
-
-      // wait for `timeout` ms and then resolve to `false` (if not cancelled first)
-      const capturedSetTimeout = setTimeout(() => {
-        if (timeout && timeout > 0) {
-          resolve(false);
-        }
-      }, timeout);
-
-      // if all promises resolve in time, cancel the timer and resolve to `true`
-      buffer.forEach(item => {
-        void resolvedSyncPromise(item).then(() => {
-          if (!--counter) {
-            clearTimeout(capturedSetTimeout);
-            resolve(true);
-          }
-        }, reject);
-      });
-    });
+    if (!buffer.size) {
+      return resolvedSyncPromise(true);
+    }
+
+    const drainPromise = drainNextSyncPromise();
+
+    if (!timeout) {
+      return drainPromise;
+    }
+
+    const promises = [drainPromise, new Promise<boolean>(resolve => setTimeout(() => resolve(false), timeout))];
+
+    // Promise.race will resolve to the first promise that resolves or rejects
+    // So if the drainPromise resolves, the timeout promise will be ignored
+    return Promise.race(promises);
   }
 
   return {
-    $: buffer,
+    get $(): PromiseLike<T>[] {
+      return Array.from(buffer);
+    },
     add,
     drain,
   };

packages/core/test/lib/utils/promisebuffer.test.ts

@@ -1,24 +1,33 @@
 import { describe, expect, test, vi } from 'vitest';
 import { makePromiseBuffer } from '../../../src/utils/promisebuffer';
-import { SyncPromise } from '../../../src/utils/syncpromise';
+import { rejectedSyncPromise, resolvedSyncPromise } from '../../../src/utils/syncpromise';
 
 describe('PromiseBuffer', () => {
   describe('add()', () => {
-    test('no limit', () => {
-      const buffer = makePromiseBuffer();
-      const p = vi.fn(() => new SyncPromise(resolve => setTimeout(resolve)));
-      void buffer.add(p);
-      expect(buffer.$.length).toEqual(1);
+    test('sync promises', () => {
+      const buffer = makePromiseBuffer(1);
+      let task1;
+      const producer1 = vi.fn(() => {
+        task1 = resolvedSyncPromise();
+        return task1;
+      });
+      const producer2 = vi.fn(() => resolvedSyncPromise());
+      expect(buffer.add(producer1)).toEqual(task1);
+      void expect(buffer.add(producer2)).rejects.toThrowError();
+      // This is immediately executed and removed again from the buffer
+      expect(buffer.$.length).toEqual(0);
+      expect(producer1).toHaveBeenCalled();
+      expect(producer2).toHaveBeenCalled();
     });
 
-    test('with limit', () => {
+    test('async promises', () => {
       const buffer = makePromiseBuffer(1);
       let task1;
       const producer1 = vi.fn(() => {
-        task1 = new SyncPromise(resolve => setTimeout(resolve));
+        task1 = new Promise(resolve => setTimeout(resolve, 1));
         return task1;
       });
-      const producer2 = vi.fn(() => new SyncPromise(resolve => setTimeout(resolve)));
+      const producer2 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
       expect(buffer.add(producer1)).toEqual(task1);
       void expect(buffer.add(producer2)).rejects.toThrowError();
       expect(buffer.$.length).toEqual(1);
@@ -28,25 +37,60 @@ describe('PromiseBuffer', () => {
   });
 
   describe('drain()', () => {
-    test('without timeout', async () => {
+    test('drains all promises without timeout', async () => {
       const buffer = makePromiseBuffer();
-      for (let i = 0; i < 5; i++) {
-        void buffer.add(() => new SyncPromise(resolve => setTimeout(resolve)));
-      }
+
+      const p1 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
+      const p2 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
+      const p3 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
+      const p4 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
+      const p5 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 1)));
+
+      [p1, p2, p3, p4, p5].forEach(p => {
+        void buffer.add(p);
+      });
+
       expect(buffer.$.length).toEqual(5);
       const result = await buffer.drain();
       expect(result).toEqual(true);
       expect(buffer.$.length).toEqual(0);
+
+      expect(p1).toHaveBeenCalled();
+      expect(p2).toHaveBeenCalled();
+      expect(p3).toHaveBeenCalled();
+      expect(p4).toHaveBeenCalled();
+      expect(p5).toHaveBeenCalled();
     });
 
-    test('with timeout', async () => {
+    test('drains all promises with timeout xxx', async () => {
       const buffer = makePromiseBuffer();
-      for (let i = 0; i < 5; i++) {
-        void buffer.add(() => new SyncPromise(resolve => setTimeout(resolve, 100)));
-      }
+
+      const p1 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 2)));
+      const p2 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 4)));
+      const p3 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 6)));
+      const p4 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 8)));
+      const p5 = vi.fn(() => new Promise(resolve => setTimeout(resolve, 10)));
+
+      [p1, p2, p3, p4, p5].forEach(p => {
+        void buffer.add(p);
+      });
+
+      expect(p1).toHaveBeenCalled();
+      expect(p2).toHaveBeenCalled();
+      expect(p3).toHaveBeenCalled();
+      expect(p4).toHaveBeenCalled();
+      expect(p5).toHaveBeenCalled();
+
       expect(buffer.$.length).toEqual(5);
-      const result = await buffer.drain(50);
+      const result = await buffer.drain(8);
       expect(result).toEqual(false);
+      // p5 is still in the buffer
+      expect(buffer.$.length).toEqual(1);
+
+      // Now drain final item
+      const result2 = await buffer.drain();
+      expect(result2).toEqual(true);
+      expect(buffer.$.length).toEqual(0);
     });
 
     test('on empty buffer', async () => {
@@ -60,7 +104,7 @@ describe('PromiseBuffer', () => {
 
   test('resolved promises should not show up in buffer length', async () => {
     const buffer = makePromiseBuffer();
-    const producer = () => new SyncPromise(resolve => setTimeout(resolve));
+    const producer = () => new Promise(resolve => setTimeout(resolve, 1));
     const task = buffer.add(producer);
     expect(buffer.$.length).toEqual(1);
     await task;
@@ -69,20 +113,18 @@ describe('PromiseBuffer', () => {
 
   test('rejected promises should not show up in buffer length', async () => {
     const buffer = makePromiseBuffer();
-    const producer = () => new SyncPromise((_, reject) => setTimeout(reject));
+    const error = new Error('whoops');
+    const producer = () => new Promise((_, reject) => setTimeout(() => reject(error), 1));
     const task = buffer.add(producer);
     expect(buffer.$.length).toEqual(1);
-    try {
-      await task;
-    } catch {
-      // no-empty
-    }
+
+    await expect(task).rejects.toThrow(error);
     expect(buffer.$.length).toEqual(0);
   });
 
   test('resolved task should give an access to the return value', async () => {
     const buffer = makePromiseBuffer<string>();
-    const producer = () => new SyncPromise<string>(resolve => setTimeout(() => resolve('test')));
+    const producer = () => resolvedSyncPromise('test');
     const task = buffer.add(producer);
     const result = await task;
     expect(result).toEqual('test');
@@ -91,7 +133,7 @@ describe('PromiseBuffer', () => {
   test('rejected task should give an access to the return value', async () => {
     expect.assertions(1);
     const buffer = makePromiseBuffer<string>();
-    const producer = () => new SyncPromise<string>((_, reject) => setTimeout(() => reject(new Error('whoops'))));
+    const producer = () => rejectedSyncPromise(new Error('whoops'));
     const task = buffer.add(producer);
     try {
       await task;