feat: add maxSize to queues (#2742)

Also restore accidentally deleted tests

Author: achingbrain

packages/utils/package.json

@@ -171,6 +171,7 @@
     "uint8arrays": "^5.1.0"
   },
   "devDependencies": {
+    "@libp2p/peer-id": "^5.0.4",
     "@types/netmask": "^2.0.5",
     "aegir": "^44.0.1",
     "benchmark": "^2.1.4",

packages/utils/src/errors.ts

@@ -18,3 +18,12 @@ export class RateLimitError extends Error {
     this.isFirstInDuration = props.isFirstInDuration
   }
 }
+
+export class QueueFullError extends Error {
+  static name = 'QueueFullError'
+
+  constructor (message: string = 'The queue was full') {
+    super(message)
+    this.name = 'QueueFullError'
+  }
+}

packages/utils/src/queue/index.ts

@@ -1,6 +1,7 @@
 import { AbortError, TypedEventEmitter } from '@libp2p/interface'
 import { pushable } from 'it-pushable'
 import { raceEvent } from 'race-event'
+import { QueueFullError } from '../errors.js'
 import { Job } from './job.js'
 import type { AbortOptions, Metrics } from '@libp2p/interface'
 
@@ -21,6 +22,14 @@ export interface QueueInit<JobReturnType, JobOptions extends AbortOptions = Abor
    */
   concurrency?: number
 
+  /**
+   * If the queue size grows to larger than this number the promise returned
+   * from the add function will reject
+   *
+   * @default Infinity
+   */
+  maxSize?: number
+
   /**
    * The name of the metric for the queue length
    */
@@ -114,6 +123,7 @@ export interface QueueEvents<JobReturnType, JobOptions extends AbortOptions = Ab
  */
 export class Queue<JobReturnType = unknown, JobOptions extends AbortOptions = AbortOptions> extends TypedEventEmitter<QueueEvents<JobReturnType, JobOptions>> {
   public concurrency: number
+  public maxSize: number
   public queue: Array<Job<JobOptions, JobReturnType>>
   private pending: number
   private readonly sort?: Comparator<Job<JobOptions, JobReturnType>>
@@ -122,6 +132,7 @@ export class Queue<JobReturnType = unknown, JobOptions extends AbortOptions = Ab
     super()
 
     this.concurrency = init.concurrency ?? Number.POSITIVE_INFINITY
+    this.maxSize = init.maxSize ?? Number.POSITIVE_INFINITY
     this.pending = 0
 
     if (init.metricName != null) {
@@ -212,6 +223,10 @@ export class Queue<JobReturnType = unknown, JobOptions extends AbortOptions = Ab
   async add (fn: RunFunction<JobOptions, JobReturnType>, options?: JobOptions): Promise<JobReturnType> {
     options?.signal?.throwIfAborted()
 
+    if (this.size === this.maxSize) {
+      throw new QueueFullError()
+    }
+
     const job = new Job<JobOptions, JobReturnType>(fn, options)
     this.enqueue(job)
     this.safeDispatchEvent('add')

packages/utils/test/peer-job-queue.spec.ts

@@ -0,0 +1,172 @@
+/* eslint-env mocha */
+
+import { generateKeyPair } from '@libp2p/crypto/keys'
+import { peerIdFromPrivateKey } from '@libp2p/peer-id'
+import { expect } from 'aegir/chai'
+import delay from 'delay'
+import pDefer from 'p-defer'
+import { raceEvent } from 'race-event'
+import { PeerQueue, type PeerQueueJobOptions } from '../src/peer-queue.js'
+import type { QueueJobFailure, QueueJobSuccess } from '../src/queue/index.js'
+
+describe('peer queue', () => {
+  it('should have jobs', async () => {
+    const deferred = pDefer()
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const privateKeyB = await generateKeyPair('Ed25519')
+    const peerIdB = peerIdFromPrivateKey(privateKeyB)
+    const queue = new PeerQueue({
+      concurrency: 1
+    })
+
+    expect(queue.has(peerIdA)).to.be.false()
+
+    void queue.add(async () => {
+      await deferred.promise
+    }, {
+      peerId: peerIdB
+    })
+
+    void queue.add(async () => {
+      await deferred.promise
+    }, {
+      peerId: peerIdA
+    })
+
+    expect(queue.has(peerIdA)).to.be.true()
+
+    deferred.resolve()
+
+    await queue.onIdle()
+
+    expect(queue.has(peerIdA)).to.be.false()
+  })
+
+  it('can join existing jobs', async () => {
+    const value = 'hello world'
+    const deferred = pDefer<string>()
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const queue = new PeerQueue<string>({
+      concurrency: 1
+    })
+
+    expect(queue.has(peerIdA)).to.be.false()
+    expect(queue.find(peerIdA)).to.be.undefined()
+
+    void queue.add(async () => {
+      return deferred.promise
+    }, {
+      peerId: peerIdA
+    })
+
+    const job = queue.find(peerIdA)
+    const join = job?.join()
+
+    deferred.resolve(value)
+
+    await expect(join).to.eventually.equal(value)
+
+    expect(queue.has(peerIdA)).to.be.false()
+    expect(queue.find(peerIdA)).to.be.undefined()
+  })
+
+  it('can join an existing job that fails', async () => {
+    const error = new Error('nope!')
+    const deferred = pDefer<string>()
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const queue = new PeerQueue<string>({
+      concurrency: 1
+    })
+
+    void queue.add(async () => {
+      return deferred.promise
+    }, {
+      peerId: peerIdA
+    })
+      .catch(() => {})
+
+    const job = queue.find(peerIdA)
+    const joinedJob = job?.join()
+
+    deferred.reject(error)
+
+    await expect(joinedJob).to.eventually.rejected
+      .with.property('message', error.message)
+  })
+
+  it('cannot join jobs after clear', async () => {
+    const value = 'hello world'
+    const deferred = pDefer<string>()
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const queue = new PeerQueue<string>({
+      concurrency: 1
+    })
+
+    expect(queue.has(peerIdA)).to.be.false()
+    expect(queue.find(peerIdA)).to.be.undefined()
+
+    void queue.add(async () => {
+      return deferred.promise
+    }, {
+      peerId: peerIdA
+    }).catch(() => {})
+
+    queue.clear()
+
+    expect(queue.find(peerIdA)).to.be.undefined()
+
+    deferred.resolve(value)
+  })
+
+  it('emits success event', async () => {
+    const value = 'hello world'
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const queue = new PeerQueue<string>({
+      concurrency: 1
+    })
+
+    void queue.add(async () => {
+      await delay(100)
+      return value
+    }, {
+      peerId: peerIdA
+    }).catch(() => {})
+
+    const event = await raceEvent<CustomEvent<QueueJobSuccess<string, PeerQueueJobOptions>>>(queue, 'success')
+
+    expect(event.detail.job.options.peerId).to.equal(peerIdA)
+    expect(event.detail.result).to.equal(value)
+  })
+
+  it('emits failure event', async () => {
+    const err = new Error('Oh no!')
+
+    const privateKeyA = await generateKeyPair('Ed25519')
+    const peerIdA = peerIdFromPrivateKey(privateKeyA)
+    const queue = new PeerQueue<string>({
+      concurrency: 1
+    })
+
+    void queue.add(async () => {
+      await delay(100)
+      throw err
+    }, {
+      peerId: peerIdA
+    }).catch(() => {})
+
+    const event = await raceEvent<CustomEvent<QueueJobFailure<string, PeerQueueJobOptions>>>(queue, 'failure')
+
+    expect(event.detail.job.options.peerId).to.equal(peerIdA)
+    expect(event.detail.error).to.equal(err)
+  })
+})

packages/utils/test/queue.spec.ts

@@ -800,4 +800,24 @@ describe('queue', () => {
     // job not in queue any more
     expect(queue.queue.find(job => !job.options.slow)).to.be.undefined()
   })
+
+  it('rejects job when the queue is full', async () => {
+    const queue = new Queue<string>({
+      concurrency: 1,
+      maxSize: 1
+    })
+
+    const job = async (): Promise<string> => {
+      await delay(100)
+      return 'hello'
+    }
+
+    const p = queue.add(job)
+
+    await expect(queue.add(job)).to.eventually.be.rejected
+      .with.property('name', 'QueueFullError')
+
+    await expect(p).to.eventually.equal('hello')
+    await expect(queue.add(job)).to.eventually.equal('hello')
+  })
 })