feat: enhance PQueue to maintain unconsumed messages after consumer restart

Signed-off-by: Vladislav Polyakov <[email protected]>

Author: polRk

packages/topic/src/queue.test.ts

@@ -52,4 +52,20 @@ describe('PQueue', () => {
 
 		await expect(queue.shift()).rejects.toThrow('Queue closed');
 	});
+
+	it('should keep unconsumed messages after consumer restart', async () => {
+		const queue = new PQueue<number>();
+
+		const shiftPromise = queue.shift();
+		queue.restartConsumer();
+		queue.push(1, 1);
+		queue.push(2, 2);
+
+		await expect(shiftPromise).rejects.toThrow('Consumer restarted');
+
+		expect(queue.size).toBe(2);
+		expect(await queue.shift()).toBe(2);
+		expect(await queue.shift()).toBe(1);
+		expect(queue.size).toBe(0);
+	});
 });

packages/topic/src/queue.ts

@@ -1,64 +1,99 @@
-export class PQueue<T> implements AsyncIterable<T>, Disposable {
-	private queue: { value: T; priority: number }[] = [];
+export class PQueue<T> {
+	private heap: { value: T; priority: number }[] = [];
 	private closed = false;
-	private resolvers: ((value: T) => void)[] = [];
+	private pendingShifts: ((value: IteratorResult<T>) => void)[] = [];
+	private pendingRejects: ((reason?: any) => void)[] = [];
 
-	push(value: T, priority: number = 0) {
-		if (this.closed) throw new Error("Queue closed");
-
-		let item = { value, priority };
-		let inserted = false;
+	get size(): number {
+		return this.heap.length;
+	}
 
-		for (let i = 0; i < this.queue.length; i++) {
-			if (this.queue[i].priority < priority) {
-				this.queue.splice(i, 0, item);
-				inserted = true;
-				break;
-			}
+	push(value: T, priority: number = 0) {
+		if (this.closed) {
+			throw new Error('Queue closed');
 		}
 
-		if (!inserted) {
-			this.queue.push(item);
+		let left = 0;
+		let right = this.heap.length;
+		while (left < right) {
+			let mid = (left + right) >> 1;
+			if (this.heap[mid].priority < priority) {
+				right = mid;
+			} else {
+				left = mid + 1;
+			}
 		}
+		this.heap.splice(left, 0, { value, priority });
 
-		if (this.resolvers.length) {
-			const nextItem = this.queue.shift()!;
-			this.resolvers.shift()?.(nextItem.value);
+		if (this.pendingShifts.length > 0) {
+			this.pendingRejects.shift()!;
+			let next = this.heap.shift()!;
+			let resolve = this.pendingShifts.shift()!;
+			resolve({ value: next.value, done: false });
 		}
 	}
 
 	async shift(): Promise<T> {
-		if (this.queue.length) {
-			return this.queue.shift()!.value;
+		if (this.closed && this.heap.length === 0) {
+			throw new Error('Queue closed');
+		}
+
+		if (this.heap.length > 0) {
+			let next = this.heap.shift()!;
+			return next.value;
 		}
-		if (this.closed) throw new Error("Queue closed");
-		return new Promise<T>((resolve) => {
-			this.resolvers.push(resolve);
+
+		return new Promise<T>((resolve, reject) => {
+			this.pendingRejects.push(reject);
+			this.pendingShifts.push(({ value, done }) => {
+				if (done) {
+					reject(new Error('Queue closed'));
+				} else {
+					resolve(value);
+				}
+			});
 		});
 	}
 
 	close() {
 		this.closed = true;
-		while (this.resolvers.length) {
-			this.resolvers.shift()?.(undefined as any);
+		while (this.pendingShifts.length > 0) {
+			this.pendingRejects.shift()!;
+			let resolve = this.pendingShifts.shift()!;
+			resolve({ value: undefined as any, done: true });
 		}
 	}
 
-	async next(): Promise<IteratorResult<T>> {
-		try {
-			return { value: await this.shift(), done: false };
-		} catch {
-			return { value: undefined, done: true };
+	restartConsumer() {
+		while (this.pendingShifts.length > 0) {
+			this.pendingShifts.shift()!;
+			let reject = this.pendingRejects.shift()!;
+			reject(new Error('Consumer restarted'));
 		}
 	}
 
-	[Symbol.asyncIterator]() {
-		return this;
+	async next(): Promise<IteratorResult<T>> {
+		if (this.closed && this.heap.length === 0) {
+			return { value: undefined as any, done: true };
+		}
+
+		if (this.heap.length > 0) {
+			let next = this.heap.shift()!;
+			return { value: next.value, done: false };
+		}
+
+		return new Promise<IteratorResult<T>>((resolve, reject) => {
+			this.pendingShifts.push(resolve);
+			this.pendingRejects.push(reject);
+		});
 	}
 
-	[Symbol.dispose]() {
-		this.close();
-		this.queue.length = 0;
-		this.resolvers.length = 0;
+	async *[Symbol.asyncIterator]() {
+		while (true) {
+			// eslint-disable-next-line no-await-in-loop
+			let { value, done } = await this.next();
+			if (done) break;
+			yield value;
+		}
 	}
 }