diff --git a/src/__tests__/generic.test.ts b/src/__tests__/generic.test.ts index 8eb0b1d..51a5b8a 100644 --- a/src/__tests__/generic.test.ts +++ b/src/__tests__/generic.test.ts @@ -1,12 +1,57 @@ -import { describe, it, expect } from "vitest"; +import { afterEach, describe, it, expect, vi } from "vitest"; import { createResumableStreamContext } from "../generic"; +import { resumeStream } from "../runtime"; +import type { Publisher, Subscriber } from "../types"; import { createInMemoryPubSubForTesting } from "../../testing-utils/in-memory-pubsub"; import { streamToBuffer, createTestingStream } from "../../testing-utils/testing-stream"; +function createResumeStreamTestContext( + get: Publisher["get"] = async () => "1", + autoAcknowledge = true +) { + let onMessage: ((message: string) => void) | undefined; + const unsubscribe = vi.fn(async () => {}); + const subscriber: Subscriber = { + connect: async () => {}, + subscribe: async (_channel, callback) => { + onMessage = callback; + }, + unsubscribe, + }; + const publisher: Publisher = { + connect: async () => {}, + publish: async () => { + if (autoAcknowledge) { + onMessage?.(""); + } + return 1; + }, + set: async () => "OK", + get, + incr: async () => 1, + }; + + return { + ctx: { + keyPrefix: "test-resume", + waitUntil: () => {}, + subscriber, + publisher, + doneWatchdogIntervalMs: 10, + } satisfies Parameters[0], + unsubscribe, + acknowledge: () => onMessage?.(""), + }; +} + describe("generic interface", () => { + afterEach(() => { + vi.useRealTimers(); + }); + it("should work with custom publisher/subscriber implementations", async () => { const { publisher, subscriber } = createInMemoryPubSubForTesting(); - + const ctx = createResumableStreamContext({ waitUntil: null, publisher, @@ -16,7 +61,7 @@ describe("generic interface", () => { const { readable, writer } = createTestingStream(); const stream = await ctx.resumableStream("test-stream", () => readable); - + writer.write("Hello "); writer.write("World!"); writer.close(); @@ -28,7 +73,7 @@ describe("generic interface", () => { it("should resume streams with custom implementations", async () => { const { publisher, subscriber } = createInMemoryPubSubForTesting(); - + const ctx = createResumableStreamContext({ waitUntil: null, publisher, @@ -39,7 +84,7 @@ describe("generic interface", () => { const { readable, writer } = createTestingStream(); // Create initial stream const stream1 = await ctx.resumableStream("test-stream-2", () => readable); - + // Resume the same stream immediately const stream2 = await ctx.resumableStream("test-stream-2", () => { throw new Error("Should not be called"); @@ -51,7 +96,7 @@ describe("generic interface", () => { expect(stream1).not.toBeNull(); expect(stream2).not.toBeNull(); - + const result1 = await streamToBuffer(stream1!); const result2 = await streamToBuffer(stream2!); expect(result1).toBe("Part 1 Part 2"); @@ -60,7 +105,7 @@ describe("generic interface", () => { it("should return null if stream is done", async () => { const { publisher, subscriber } = createInMemoryPubSubForTesting(); - + const ctx = createResumableStreamContext({ waitUntil: null, publisher, @@ -70,20 +115,126 @@ describe("generic interface", () => { const { readable, writer } = createTestingStream(); const stream = await ctx.resumableStream("test-stream-3", () => readable); - + writer.write("Done"); writer.close(); - + await streamToBuffer(stream!); - + // Try to resume after stream is done const doneStream = await ctx.resumableStream("test-stream-3", () => { throw new Error("Should not be called"); }); - + expect(doneStream).toBeNull(); }); + it( + "closes resumed streams via the done watchdog when the DONE pub/sub message is lost", + { timeout: 5000 }, + async () => { + const { publisher, subscriber } = createInMemoryPubSubForTesting(); + + // Pub/sub is fire-and-forget; simulate the DONE control message getting lost in transit + // (e.g. subscriber reconnect) while durable writes (the sentinel SET) still go through. + const lossyPublisher: typeof publisher = { + ...publisher, + publish: async (channel: string, message: string) => { + if (message.includes("DONE_SENTINEL")) { + return 0; + } + return publisher.publish(channel, message); + }, + }; + + const ctx = createResumableStreamContext({ + waitUntil: null, + publisher: lossyPublisher, + subscriber, + keyPrefix: "test-generic-" + crypto.randomUUID(), + doneWatchdogIntervalMs: 25, + }); + + const { readable, writer } = createTestingStream(); + const producerStream = await ctx.resumableStream("test-stream-watchdog", () => readable); + const resumedStream = await ctx.resumableStream("test-stream-watchdog", () => { + throw new Error("Should not be called"); + }); + + writer.write("Hello "); + writer.write("World!"); + writer.close(); + + expect(await streamToBuffer(producerStream!)).toBe("Hello World!"); + // Without the watchdog this hangs forever: the consumer never receives the DONE message. + expect(await streamToBuffer(resumedStream!)).toBe("Hello World!"); + } + ); + + it("cleans up the done watchdog when subscription setup fails", async () => { + vi.useFakeTimers(); + const get = vi.fn(async () => "1"); + const { ctx, unsubscribe } = createResumeStreamTestContext(get); + ctx.subscriber.subscribe = async () => { + throw new Error("subscribe failed"); + }; + + await expect(resumeStream(ctx, "failed-subscription")).rejects.toThrow("subscribe failed"); + await vi.advanceTimersByTimeAsync(1000); + + expect(get).not.toHaveBeenCalled(); + expect(unsubscribe).toHaveBeenCalledOnce(); + }); + + it("cleans up the done watchdog when the resumed stream is cancelled", async () => { + vi.useFakeTimers(); + const get = vi.fn(async () => "1"); + const { ctx, unsubscribe } = createResumeStreamTestContext(get); + const stream = await resumeStream(ctx, "cancelled-stream"); + + await stream!.cancel(); + await vi.advanceTimersByTimeAsync(1000); + + expect(get).not.toHaveBeenCalled(); + expect(unsubscribe).toHaveBeenCalledOnce(); + }); + + it("does not start the done watchdog before the initial ack", async () => { + vi.useFakeTimers(); + const get = vi.fn(async () => "1"); + const { ctx, acknowledge } = createResumeStreamTestContext(get, false); + const streamPromise = resumeStream(ctx, "delayed-ack"); + + await vi.advanceTimersByTimeAsync(100); + expect(get).not.toHaveBeenCalled(); + + acknowledge(); + const stream = await streamPromise; + await stream!.cancel(); + }); + + it("does not overlap done watchdog checks when Redis is slow", async () => { + vi.useFakeTimers(); + let resolveGet: (value: string) => void; + const get = vi.fn(() => { + return new Promise((resolve) => { + resolveGet = resolve; + }); + }); + const { ctx } = createResumeStreamTestContext(get); + const stream = await resumeStream(ctx, "slow-redis"); + + vi.advanceTimersByTime(10); + expect(get).toHaveBeenCalledOnce(); + + vi.advanceTimersByTime(35); + expect(get).toHaveBeenCalledOnce(); + + resolveGet("1"); + await Promise.resolve(); + await stream!.cancel(); + }); + it("should throw error if publisher is not provided", () => { expect(() => { createResumableStreamContext({ @@ -93,4 +244,3 @@ describe("generic interface", () => { }).toThrow(); }); }); - diff --git a/src/runtime.ts b/src/runtime.ts index b4dd3fe..b366d92 100644 --- a/src/runtime.ts +++ b/src/runtime.ts @@ -9,6 +9,7 @@ interface CreateResumableStreamContext { waitUntil: (promise: Promise) => void; subscriber: Subscriber; publisher: Publisher; + doneWatchdogIntervalMs?: number; } export function createResumableStreamContextFactory(defaults: _Private.RedisDefaults) { @@ -21,6 +22,7 @@ export function createResumableStreamContextFactory(defaults: _Private.RedisDefa waitUntil, subscriber: options.subscriber, publisher: options.publisher, + doneWatchdogIntervalMs: options.doneWatchdogIntervalMs, } as CreateResumableStreamContext; let initPromises: Promise[] = []; @@ -107,6 +109,8 @@ const DONE_MESSAGE = "\n\n\nDONE_SENTINEL_hasdfasudfyge374%$%^$EDSATRTYFtydryrte const DONE_VALUE = "DONE"; +const DEFAULT_DONE_WATCHDOG_INTERVAL_MS = 10_000; + async function resumeExistingStream( initPromise: Promise, ctx: CreateResumableStreamContext, @@ -254,17 +258,94 @@ export async function resumeStream( ): Promise | null> { const listenerId = crypto.randomUUID(); return new Promise | null>((resolve, reject) => { + const chunkChannel = `${ctx.keyPrefix}:chunk:${listenerId}`; + const sentinelKey = `${ctx.keyPrefix}:sentinel:${streamId}`; + const watchdogIntervalMs = ctx.doneWatchdogIntervalMs ?? DEFAULT_DONE_WATCHDOG_INTERVAL_MS; + let ackTimeout: ReturnType | undefined; + let watchdogTimeout: ReturnType | undefined; + let cleanupPromise: Promise | undefined; + + const cleanup = () => { + if (cleanupPromise) { + return cleanupPromise; + } + + clearTimeout(ackTimeout); + clearTimeout(watchdogTimeout); + cleanupPromise = Promise.resolve().then(() => ctx.subscriber.unsubscribe(chunkChannel)); + return cleanupPromise; + }; + const readableStream = new ReadableStream({ async start(controller) { try { debugLog("STARTING STREAM", streamId, listenerId); - const cleanup = async () => { - await ctx.subscriber.unsubscribe(`${ctx.keyPrefix}:chunk:${listenerId}`); + // The DONE control message travels over pub/sub, which is fire-and-forget: if that one + // message is lost (e.g. the subscriber connection dropped and reconnected at the wrong + // moment, or the producer died between writing the DONE sentinel and publishing), this + // stream would stay open forever even though the durable sentinel already says DONE. + // Re-check the sentinel periodically and close once the producer is finished or the + // sentinel expired. Closing requires two consecutive DONE observations so in-flight + // messages get a full interval to drain before we give up on them. + let doneObservations = 0; + let watchdogStarted = false; + const closeStream = () => { + try { + controller.close(); + } catch (e) { + // The stream may already be closed because the client disconnected. + if (isDebug()) { + console.error(e); + } + } }; + const scheduleDoneWatchdog = () => { + if (cleanupPromise) { + return; + } + watchdogTimeout = setTimeout(checkDone, watchdogIntervalMs); + }; + const startDoneWatchdog = () => { + if (watchdogStarted) { + return; + } + watchdogStarted = true; + scheduleDoneWatchdog(); + }; + async function checkDone() { + try { + const val = await ctx.publisher.get(sentinelKey); + if (val !== DONE_VALUE && val !== null) { + doneObservations = 0; + return; + } + + doneObservations += 1; + if (doneObservations < 2) { + return; + } + + debugLog( + "done watchdog: sentinel is done but no DONE message arrived; closing", + streamId, + listenerId + ); + closeStream(); + await cleanup(); + } catch (e) { + // A transient sentinel read failure must not kill the stream. Cleanup failures are + // also contained here because timer callbacks cannot surface rejected promises. + if (isDebug()) { + console.error(e); + } + } finally { + scheduleDoneWatchdog(); + } + } const start = Date.now(); - const timeout = setTimeout(async () => { + ackTimeout = setTimeout(async () => { await cleanup(); - const val = await ctx.publisher.get(`${ctx.keyPrefix}:sentinel:${streamId}`); + const val = await ctx.publisher.get(sentinelKey); if (val === DONE_VALUE) { resolve(null); } @@ -273,40 +354,29 @@ export async function resumeStream( reject(new Error("Timeout waiting for ack")); } }, 1000); - await ctx.subscriber.subscribe( - `${ctx.keyPrefix}:chunk:${listenerId}`, - async (message: string) => { - debugLog("Received message", message); - // The other side always sends a message even if it is the empty string. - clearTimeout(timeout); - resolve(readableStream); - if (message === DONE_MESSAGE) { - try { - controller.close(); - } catch (e) { - // errors can e.g. happen if the stream is already closed - // because the client has disconnected - // ignore them unless we are in debug mode - if (isDebug()) { - console.error(e); - } - } - await cleanup(); - return; - } - try { - controller.enqueue(message); - } catch (e) { - // errors can e.g. happen if the stream is already closed - // because the client has disconnected - // ignore them unless we are in debug mode - if (isDebug()) { - console.error(e); - } - await cleanup(); + await ctx.subscriber.subscribe(chunkChannel, async (message: string) => { + debugLog("Received message", message); + // The other side always sends a message even if it is the empty string. + clearTimeout(ackTimeout); + resolve(readableStream); + if (message === DONE_MESSAGE) { + closeStream(); + await cleanup(); + return; + } + startDoneWatchdog(); + try { + controller.enqueue(message); + } catch (e) { + // errors can e.g. happen if the stream is already closed + // because the client has disconnected + // ignore them unless we are in debug mode + if (isDebug()) { + console.error(e); } + await cleanup(); } - ); + }); await ctx.publisher.publish( `${ctx.keyPrefix}:request:${streamId}`, JSON.stringify({ @@ -315,9 +385,19 @@ export async function resumeStream( }) ); } catch (e) { + try { + await cleanup(); + } catch (cleanupError) { + if (isDebug()) { + console.error(cleanupError); + } + } reject(e); } }, + async cancel() { + await cleanup(); + }, }); }); } diff --git a/src/types.ts b/src/types.ts index 09d6390..d24b00d 100644 --- a/src/types.ts +++ b/src/types.ts @@ -21,6 +21,12 @@ export interface CreateResumableStreamContextOptions { * A pubsub publisher. Designed to be compatible with clients from the `redis` package. */ publisher?: Publisher | Redis; + /** + * How often resumed streams re-check the durable DONE sentinel as a fallback for a lost DONE + * pub/sub message. A resumed stream closes after two consecutive checks observe a finished + * (or expired) sentinel. Defaults to 10000. + */ + doneWatchdogIntervalMs?: number; } export interface ResumableStreamContext {