From 8d0a35651f6629e1a6ba1e224043a7699619902c Mon Sep 17 00:00:00 2001
From: Dobes Vandermeer <dobesv@gmail.com>
Date: Tue, 30 Mar 2021 22:26:56 -0700
Subject: [PATCH] Add `.onSizeLessThan()` helper method (#131)

Co-authored-by: Sindre Sorhus <sindresorhus@gmail.com>
---
 readme.md       |  8 ++++++++
 source/index.ts | 25 +++++++++++++++++++++++++
 test/test.ts    | 26 ++++++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/readme.md b/readme.md
index 142851f..4323c07 100644
--- a/readme.md
+++ b/readme.md
@@ -155,6 +155,14 @@ Returns a promise that settles when the queue becomes empty, and all promises ha
 
 The difference with `.onEmpty` is that `.onIdle` guarantees that all work from the queue has finished. `.onEmpty` merely signals that the queue is empty, but it could mean that some promises haven't completed yet.
 
+#### .onSizeLessThan(limit)
+
+Returns a promise that settles when the queue size is less than the given limit: `queue.size < limit`.
+
+If you want to avoid having the queue grow beyond a certain size you can `await queue.onSizeLessThan()` before adding a new item.
+
+Note that this only limits the number of items waiting to start. There could still be up to `concurrency` jobs already running that this call does not include in its calculation.
+
 #### .clear()
 
 Clear the queue.
diff --git a/source/index.ts b/source/index.ts
index 4818c8a..6397d64 100644
--- a/source/index.ts
+++ b/source/index.ts
@@ -328,6 +328,31 @@ export default class PQueue<QueueType extends Queue<RunFunction, EnqueueOptionsT
 		});
 	}
 
+	/**
+	@returns A promise that settles when the queue size is less than the given limit: `queue.size < limit`.
+
+	If you want to avoid having the queue grow beyond a certain size you can `await queue.onSizeLessThan()` before adding a new item.
+
+	Note that this only limits the number of items waiting to start. There could still be up to `concurrency` jobs already running that this call does not include in its calculation.
+	*/
+	async onSizeLessThan(limit: number): Promise<void> {
+		// Instantly resolve if the queue is empty.
+		if (this._queue.size < limit) {
+			return;
+		}
+
+		return new Promise<void>(resolve => {
+			const listener = () => {
+				if (this._queue.size < limit) {
+					this.removeListener('next', listener);
+					resolve();
+				}
+			};
+
+			this.on('next', listener);
+		});
+	}
+
 	/**
 	The difference with `.onEmpty` is that `.onIdle` guarantees that all work from the queue has finished. `.onEmpty` merely signals that the queue is empty, but it could mean that some promises haven't completed yet.
 
diff --git a/test/test.ts b/test/test.ts
index f7f2d58..80eb95c 100644
--- a/test/test.ts
+++ b/test/test.ts
@@ -225,6 +225,32 @@ test('.onIdle()', async t => {
 	t.is(queue.pending, 0);
 });
 
+test('.onSizeLessThan()', async t => {
+	const queue = new PQueue({concurrency: 1});
+
+	queue.add(async () => delay(100));
+	queue.add(async () => delay(100));
+	queue.add(async () => delay(100));
+	queue.add(async () => delay(100));
+	queue.add(async () => delay(100));
+
+	await queue.onSizeLessThan(4);
+	t.is(queue.size, 3);
+	t.is(queue.pending, 1);
+
+	await queue.onSizeLessThan(2);
+	t.is(queue.size, 1);
+	t.is(queue.pending, 1);
+
+	await queue.onSizeLessThan(10);
+	t.is(queue.size, 1);
+	t.is(queue.pending, 1);
+
+	await queue.onSizeLessThan(1);
+	t.is(queue.size, 0);
+	t.is(queue.pending, 1);
+});
+
 test('.onIdle() - no pending', async t => {
 	const queue = new PQueue();
 	t.is(queue.size, 0);