diff --git a/src/__tests__/tests.ts b/src/__tests__/tests.ts index 867e5e4..88fb412 100644 --- a/src/__tests__/tests.ts +++ b/src/__tests__/tests.ts @@ -267,5 +267,22 @@ export function resumableStreamTests( expect(await resume.resumeExistingStream("test")).toBeNull(); expect(result).toEqual("1\n2\n"); }); + + it("should finish the stream when the source errors", async () => { + const { readable, writer } = createTestingStream(); + const stream = await resume.createNewResumableStream("test", () => readable); + writer.write("1\n"); + const reader = stream!.getReader(); + expect(await reader.read()).toEqual({ done: false, value: "1\n" }); + const stream2 = await resume.resumeExistingStream("test"); + expect(await streamToBuffer(stream2, 1)).toEqual("1\n"); + + writer.abort(new Error("source failed")); + + await expect(reader.read()).rejects.toThrow("source failed"); + expect(await streamToBuffer(stream2)).toEqual(""); + expect(await resume.hasExistingStream("test")).toBe("DONE"); + expect(await resume.resumeExistingStream("test")).toBeNull(); + }); }); } diff --git a/src/runtime.ts b/src/runtime.ts index e0dfeb6..cc982b0 100644 --- a/src/runtime.ts +++ b/src/runtime.ts @@ -173,50 +173,76 @@ async function createNewResumableStream( start(controller) { const stream = makeStream(); const reader = stream.getReader(); + async function finish(error?: unknown) { + if (isDone) { + return; + } + isDone = true; + debugLog(error === undefined ? "Stream done" : "Stream errored"); + try { + if (error === undefined) { + controller.close(); + } else { + controller.error(error); + } + } catch (e) { + //console.error(e); + } + const promises: Promise[] = []; + debugLog("setting sentinel to done"); + promises.push( + ctx.publisher.set(`${ctx.keyPrefix}:sentinel:${streamId}`, DONE_VALUE, { + EX: 24 * 60 * 60, + }) + ); + promises.push(ctx.subscriber.unsubscribe(`${ctx.keyPrefix}:request:${streamId}`)); + for (const listenerId of listenerChannels) { + debugLog("sending done message to", listenerId); + promises.push( + ctx.publisher.publish(`${ctx.keyPrefix}:chunk:${listenerId}`, DONE_MESSAGE) + ); + } + await Promise.all(promises); + streamDoneResolver?.(); + debugLog("Cleanup done"); + } function read() { - reader.read().then(async ({ done, value }) => { - if (done) { - isDone = true; - debugLog("Stream done"); + reader + .read() + .then(async ({ done, value }) => { + if (done) { + await finish(); + return; + } + chunks.push(value); try { - controller.close(); + debugLog("Enqueuing line", value); + controller.enqueue(value); } catch (e) { - //console.error(e); + // If we cannot enqueue, the stream is already closed, but we WANT to continue. } const promises: Promise[] = []; - debugLog("setting sentinel to done"); - promises.push( - ctx.publisher.set(`${ctx.keyPrefix}:sentinel:${streamId}`, DONE_VALUE, { - EX: 24 * 60 * 60, - }) - ); - promises.push(ctx.subscriber.unsubscribe(`${ctx.keyPrefix}:request:${streamId}`)); for (const listenerId of listenerChannels) { - debugLog("sending done message to", listenerId); - promises.push( - ctx.publisher.publish(`${ctx.keyPrefix}:chunk:${listenerId}`, DONE_MESSAGE) - ); + debugLog("sending line to", listenerId); + promises.push(ctx.publisher.publish(`${ctx.keyPrefix}:chunk:${listenerId}`, value)); } await Promise.all(promises); - streamDoneResolver?.(); - debugLog("Cleanup done"); - return; - } - chunks.push(value); - try { - debugLog("Enqueuing line", value); - controller.enqueue(value); - } catch (e) { - // If we cannot enqueue, the stream is already closed, but we WANT to continue. - } - const promises: Promise[] = []; - for (const listenerId of listenerChannels) { - debugLog("sending line to", listenerId); - promises.push(ctx.publisher.publish(`${ctx.keyPrefix}:chunk:${listenerId}`, value)); - } - await Promise.all(promises); - read(); - }); + read(); + }) + .catch(async (error) => { + // The source stream errored, or a Redis call failed. Either way nothing reads the + // source anymore, so finish the stream instead of leaving the sentinel in progress. + reader.cancel(error).catch(() => {}); + try { + await finish(error); + } catch (e) { + if (isDebug()) { + console.error(e); + } + } finally { + streamDoneResolver?.(); + } + }); } read(); },