diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx index 179c94c7fe28..49b3cf2451dc 100644 --- a/packages/next/server/render.tsx +++ b/packages/next/server/render.tsx @@ -1250,7 +1250,8 @@ export async function renderToHTML( ReactDOMServer, content, suffix, - serverComponentsInlinedTransformStream?.readable ?? null, + serverComponentsInlinedTransformStream?.readable ?? + streamFromArray([]), generateStaticHTML ) } @@ -1382,7 +1383,7 @@ export async function renderToHTML( ReactDOMServer, document, null, - null, + streamFromArray([]), true ) documentHTML = await streamToString(documentStream) @@ -1523,8 +1524,8 @@ function createTransformStream({ flush, transform, }: { - flush: (controller: TransformStreamDefaultController) => Promise | void - transform: ( + flush?: (controller: TransformStreamDefaultController) => Promise | void + transform?: ( chunk: Uint8Array, controller: TransformStreamDefaultController ) => void @@ -1560,7 +1561,7 @@ function createTransformStream({ const { done, value } = await reader.read() if (done) { - const maybePromise = flush(controller) + const maybePromise = flush?.(controller) if (maybePromise) { await maybePromise } @@ -1568,7 +1569,11 @@ function createTransformStream({ return } - transform(value, controller) + if (transform) { + transform(value, controller) + } else { + controller.enqueue(value) + } } } catch (err) { writer.abort(err) @@ -1620,7 +1625,7 @@ function renderToStream( ReactDOMServer: typeof import('react-dom/server'), element: React.ReactElement, suffix: string | null, - dataStream: ReadableStream | null, + dataStream: ReadableStream, generateStaticHTML: boolean ): Promise { return new Promise((resolve, reject) => { @@ -1635,83 +1640,109 @@ function renderToStream( } : null - const { readable, writable } = createTransformStream({ - transform(chunk, controller) { - controller.enqueue(chunk) - }, - - flush(controller) { - if (suffixState) { - controller.enqueue(suffixState.closeTag) - } - }, - }) - const doResolve = () => { if (!resolved) { resolved = true - let dataStreamFinished: Promise | null = null - let shellFlushed = false - - resolve( - readable.pipeThrough(createBufferedTransformStream()).pipeThrough( - createTransformStream({ - transform(chunk, controller) { - controller.enqueue(chunk) - - if (!shellFlushed) { - shellFlushed = true - if (suffixState) { - controller.enqueue(suffixState.suffixUnclosed) - } - } - if (!dataStreamFinished && dataStream) { - const dataStreamReader = dataStream.getReader() - dataStreamFinished = (async () => { - try { - while (true) { - const { done, value } = await dataStreamReader.read() - if (done) { - return - } - controller.enqueue(value) - } - } catch (err) { - controller.error(err) - } - })() - } - }, - flush() { - if (dataStreamFinished) { - return dataStreamFinished - } - }, - }) + // React will call our callbacks synchronously, so we need to + // defer to a microtask to ensure `stream` is set. + Promise.resolve().then(() => + resolve( + stream + .pipeThrough(createBufferedTransformStream()) + .pipeThrough( + createInlineDataStream( + dataStream.pipeThrough( + createPrefixStream(suffixState?.suffixUnclosed ?? null) + ) + ) + ) + .pipeThrough(createSuffixStream(suffixState?.closeTag ?? null)) ) ) } } - ;(ReactDOMServer as any) - .renderToReadableStream(element, { - onError(err: Error) { - if (!resolved) { - resolved = true - reject(err) - } - }, - onCompleteShell() { - if (!generateStaticHTML) { - doResolve() - } - }, - onCompleteAll() { + const stream: ReadableStream = ( + ReactDOMServer as any + ).renderToReadableStream(element, { + onError(err: Error) { + if (!resolved) { + resolved = true + reject(err) + } + }, + onCompleteShell() { + if (!generateStaticHTML) { doResolve() - }, - }) - .pipeTo(writable) + } + }, + onCompleteAll() { + doResolve() + }, + }) + }) +} + +function createSuffixStream(suffix: Uint8Array | null) { + return createTransformStream({ + flush(controller) { + if (suffix) { + controller.enqueue(suffix) + } + }, + }) +} + +function createPrefixStream(prefix: Uint8Array | null) { + let prefixFlushed = false + return createTransformStream({ + transform(chunk, controller) { + if (!prefixFlushed && prefix) { + prefixFlushed = true + controller.enqueue(prefix) + } + controller.enqueue(chunk) + }, + flush(controller) { + if (!prefixFlushed && prefix) { + prefixFlushed = true + controller.enqueue(prefix) + } + }, + }) +} + +function createInlineDataStream( + dataStream: ReadableStream | null +): TransformStream { + let dataStreamFinished: Promise | null = null + return createTransformStream({ + transform(chunk, controller) { + controller.enqueue(chunk) + + if (!dataStreamFinished && dataStream) { + const dataStreamReader = dataStream.getReader() + dataStreamFinished = (async () => { + try { + while (true) { + const { done, value } = await dataStreamReader.read() + if (done) { + return + } + controller.enqueue(value) + } + } catch (err) { + controller.error(err) + } + })() + } + }, + flush() { + if (dataStreamFinished) { + return dataStreamFinished + } + }, }) } diff --git a/test/integration/react-streaming-and-server-components/test/streaming.js b/test/integration/react-streaming-and-server-components/test/streaming.js index 7748400b92b8..9aedda759cf6 100644 --- a/test/integration/react-streaming-and-server-components/test/streaming.js +++ b/test/integration/react-streaming-and-server-components/test/streaming.js @@ -104,4 +104,11 @@ export default function (context) { 'count: 1' ) }) + + it('should flush the suffix at the very end', async () => { + await fetchViaHTTP(context.appPort, '/').then(async (response) => { + const result = await resolveStreamResponse(response) + expect(result).toMatch(/<\/body><\/html>/) + }) + }) }