feat: promise queue implementation (#771)

* feat: promise queue implementation

* chore: linter changes

* tests: promise and priority queue tests

* chore: linter changes

* Apply suggestions from code review

Co-authored-by: Alex Ruzenhack <[email protected]>

* chore: docstrings

* chore: review requests

* Update src/models/promise_queue.ts

Co-authored-by: André Abadesso <[email protected]>

* chore: change pop empty check

---------

Co-authored-by: Alex Ruzenhack <[email protected]>
Co-authored-by: André Abadesso <[email protected]>

Author: 3 people

__tests__/models/priority_queue.test.ts

@@ -0,0 +1,46 @@
+/**
+ * Copyright (c) Hathor Labs and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import PriorityQueue from '../../src/models/priority_queue';
+
+test('PriorityQueue operations', () => {
+  const q = new PriorityQueue<string>();
+  expect(q.size).toEqual(0);
+  // method: makeNode
+  expect(PriorityQueue.makeNode(123, 'foobar')).toMatchObject({
+    priority: 123,
+    value: 'foobar',
+  });
+
+  // methods: push and pop
+  q.push(PriorityQueue.makeNode(10, 'foobar1'));
+  expect(q.size).toEqual(1);
+  expect(q.isEmpty()).toBeFalsy();
+  expect(q.pop()).toEqual('foobar1');
+  expect(q.isEmpty()).toBeTruthy();
+
+  // method: add
+  const elements = [
+    { priority: 10, value: 'foo' },
+    { priority: 1, value: 'baz' },
+    { priority: 5, value: 'bar' },
+  ].map(el => PriorityQueue.makeNode(el.priority, el.value));
+  q.add(...elements);
+  expect(q.size).toEqual(3);
+  expect(q.isEmpty()).toBeFalsy();
+  expect(q.pop()).toEqual('foo');
+  expect(q.size).toEqual(2);
+  expect(q.isEmpty()).toBeFalsy();
+  expect(q.peek()).toEqual('bar');
+  expect(q.pop()).toEqual('bar');
+  expect(q.size).toEqual(1);
+  expect(q.isEmpty()).toBeFalsy();
+  expect(q.peek()).toEqual('baz');
+  expect(q.pop()).toEqual('baz');
+  expect(q.size).toEqual(0);
+  expect(q.isEmpty()).toBeTruthy();
+});

__tests__/models/promise_queue.test.ts

@@ -0,0 +1,55 @@
+/**
+ * Copyright (c) Hathor Labs and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import PromiseQueue from '../../src/models/promise_queue';
+
+describe('PromiseQueue operations', () => {
+  it('should start unpaused and be pausable', () => {
+    const q = new PromiseQueue();
+    expect(q.isPaused).toBeFalsy();
+    q.stop();
+    expect(q.isPaused).toBeTruthy();
+    q.continue();
+    expect(q.isPaused).toBeFalsy();
+  });
+
+  it('should be able to set the concurrency', async () => {
+    const q = new PromiseQueue();
+    expect(q.concurrent).toEqual(1);
+    q.concurrent = 3;
+    expect(q.concurrent).toEqual(3);
+
+    // Check that the concurrent actually runs only 2 tasks
+    q.concurrent = 2;
+
+    const task = async () => {
+      return new Promise<void>(resolve => {
+        setTimeout(resolve, 2000);
+      });
+    };
+    q.add(task);
+    q.add(task);
+    q.add(task);
+    // wait 50ms so jobs can start
+    await new Promise<void>(resolve => {
+      setTimeout(resolve, 50);
+    });
+    // Expect that only 2 jobs are running
+    expect(q.jobsRunning).toEqual(2);
+
+    // Wait for either all tasks to finish or a timeout.
+    await Promise.race([
+      new Promise<void>(resolve => {
+        // idle is sent when the queue  is empty and there are no events running
+        q.on('idle', resolve);
+      }),
+      new Promise<void>((_, reject) => {
+        setTimeout(reject, 5000);
+      }),
+    ]);
+  });
+});

babel.config.js

@@ -15,5 +15,6 @@ module.exports = {
   plugins: [
     "@babel/plugin-transform-async-generator-functions",
     "@babel/plugin-transform-class-properties",
+    "@babel/plugin-transform-private-methods",
   ],
 };

package-lock.json

@@ -24,7 +24,6 @@
     "level": "8.0.1",
     "lodash": "4.17.21",
     "long": "5.2.3",
-    "queue-promise": "^2.2.1",
     "ws": "8.17.1"
   },
   "scripts": {

package.json

@@ -0,0 +1,186 @@
+/**
+ * Copyright (c) Hathor Labs and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+/* eslint class-methods-use-this: ["error", { "exceptMethods": ["parent", "left", "right", "comparator"] }] */
+
+interface PriorityQueueNode<T> {
+  value: T;
+  priority: number;
+}
+
+/**
+ * Priority queue implementation using an underlying heap for performance.
+ * @example
+ * const prioQ = new PriorityQueue<string>();
+ * prioQ.push(prioQ.makeNode(0, 'lower prio'));
+ * prioQ.push(prioQ.makeNode(10, 'higher prio'));
+ * prioQ.push(prioQ.makeNode(5, 'medium prio'));
+ * // Returns 'higher prio'
+ * prioQ.pop()
+ * // Returns 'medium prio'
+ * prioQ.pop()
+ * // Returns 'lower prio'
+ * prioQ.pop()
+ *
+ * @template [T=unknown]
+ */
+export default class PriorityQueue<T = unknown> {
+  #heap: PriorityQueueNode<T>[];
+
+  readonly #top: number;
+
+  constructor() {
+    this.#heap = [];
+    this.#top = 0;
+  }
+
+  /**
+   * Utility to create a node from a value and priority.
+   * @example
+   * // returns { value: 'foobar', 'priority': 10 }
+   * prioQueue.makeNode(10, 'foobar');
+   */
+  static makeNode<N>(priority: number, value: N): PriorityQueueNode<N> {
+    return { value, priority };
+  }
+
+  /**
+   * Get the number of elements on the priority queue.
+   */
+  public get size() {
+    return this.#heap.length;
+  }
+
+  /**
+   * Check if the priority queue is empty
+   */
+  public isEmpty() {
+    return this.size === 0;
+  }
+
+  /**
+   * Get the top value of the queue without removing it from the queue.
+   */
+  public peek() {
+    if (this.size === 0) {
+      // Heap is empty
+      return undefined;
+    }
+    return this.#heap[this.#top].value;
+  }
+
+  /**
+   * Add a node to the priority queue and maintain the priority order.
+   */
+  public push(value: PriorityQueueNode<T>) {
+    this.#heap.push(value);
+    this._siftUp();
+  }
+
+  /**
+   * Add multiple values to the priority queue while maintaining the priority order.
+   */
+  public add(...nodes: PriorityQueueNode<T>[]) {
+    nodes.forEach(node => {
+      this.push(node);
+    });
+    return this.size;
+  }
+
+  /**
+   * Get the node with highest priority and remove it from the priority queue.
+   */
+  public pop() {
+    if (this.isEmpty()) {
+      // Queue is empty
+      return undefined;
+    }
+    // This may be undefined if the queue is empty, but we already checked for that.
+    const poppedValue = this.peek();
+    const bottom = this.size - 1;
+    if (bottom > this.#top) {
+      this._swap(this.#top, bottom);
+    }
+    this.#heap.pop();
+    this._siftDown();
+    return poppedValue;
+  }
+
+  /** Compare 2 nodes and return true if the left one has higher priority. */
+  private comparator(a: PriorityQueueNode<T>, b: PriorityQueueNode<T>): boolean {
+    return a.priority > b.priority;
+  }
+
+  /** Given a node index on the heap get the parent index */
+  private parent(i: number) {
+    return ((i + 1) >>> 1) - 1;
+  }
+
+  /** Given a node index on the heap get the left child index */
+  private left(i: number) {
+    return (i << 1) + 1;
+  }
+
+  /** Given a node index on the heap get the right child index */
+  private right(i: number) {
+    return (i + 1) << 1;
+  }
+
+  /**
+   * Compare the nodes at index `i`, `j` on the heap
+   * Return true if the node at `i` is higher priority than the node at `j`
+   * Return false otherwise.
+   */
+  private _greater(i: number, j: number) {
+    return this.comparator(this.#heap[i], this.#heap[j]);
+  }
+
+  /** swap the nodes at index `i` and `j` on the heap */
+  private _swap(i: number, j: number) {
+    [this.#heap[i], this.#heap[j]] = [this.#heap[j], this.#heap[i]];
+  }
+
+  /**
+   * The last node of the heap will work its way up the heap until it meets a node
+   * of higher priority or reaches the top of the heap.
+   */
+  private _siftUp() {
+    // Start from the last index and work our way up the heap
+    let node = this.size - 1;
+    // While the current node is not at the top and the priority is greater
+    // than the parent we continue sifting up
+    while (node > this.#top && this._greater(node, this.parent(node))) {
+      // parent is lower priority, swap with child (current node)
+      this._swap(node, this.parent(node));
+      // Start from the parent
+      node = this.parent(node);
+    }
+  }
+
+  /**
+   * The top node of the heap will work its way down until no child is of higher
+   * priority or it reaches the bottom of the heap.
+   */
+  private _siftDown() {
+    // Start from the top index and work our way down the heap
+    let node = this.#top;
+    // If a child is in the heap and has higher priority, swap with parent and go down
+    while (
+      (this.left(node) < this.size && this._greater(this.left(node), node)) ||
+      (this.right(node) < this.size && this._greater(this.right(node), node))
+    ) {
+      // Get the child with higher priority
+      const maxChild =
+        this.right(node) < this.size && this._greater(this.right(node), this.left(node))
+          ? this.right(node)
+          : this.left(node);
+      // Swap with parent and continue sifting down from the child
+      this._swap(node, maxChild);
+      // Start from the child
+      node = maxChild;
+    }
+  }
+}