fix: more robust stream handling (TanStack#5996)

Author: schiller-manuel

packages/react-router/src/ssr/renderRouterToStream.tsx

@@ -61,10 +61,17 @@ export const renderRouterToStream = async ({
  }),
  onError: (error, info) => {
  console.error('Error in renderToPipeableStream:', error, info)
+ // Destroy the passthrough stream on error
+ if (!reactAppPassthrough.destroyed) {
+ reactAppPassthrough.destroy(
+ error instanceof Error ? error : new Error(String(error)),
+ )
+ }
  },
  })
  } catch (e) {
  console.error('Error in renderToPipeableStream:', e)
+ reactAppPassthrough.destroy(e instanceof Error ? e : new Error(String(e)))
  }
 
  const responseStream = transformPipeableStreamWithRouter(

packages/react-router/src/ssr/renderRouterToString.tsx

@@ -28,5 +28,7 @@ export const renderRouterToString = async ({
  status: 500,
  headers: responseHeaders,
  })
+ } finally {
+ router.serverSsr?.cleanup()
  }
 }

packages/router-core/src/ssr/createRequestHandler.ts

@@ -20,40 +20,54 @@ export function createRequestHandler<TRouter extends AnyRouter>({
 }): RequestHandler<TRouter> {
  return async (cb) => {
  const router = createRouter()
+ // Track whether the callback will handle cleanup
+ let cbWillCleanup = false
 
- attachRouterServerSsrUtils({
- router,
- manifest: await getRouterManifest?.(),
- })
+ try {
+ attachRouterServerSsrUtils({
+ router,
+ manifest: await getRouterManifest?.(),
+ })
 
- const url = new URL(request.url, 'http://localhost')
- const origin = getOrigin(request)
- const href = url.href.replace(url.origin, '')
+  const url = new URL(request.url, 'http://localhost')
+  const origin = getOrigin(request)
+  const href = url.href.replace(url.origin, '')
 
- // Create a history for the router
- const history = createMemoryHistory({
- initialEntries: [href],
- })
+  // Create a history for the router
+  const history = createMemoryHistory({
+  initialEntries: [href],
+  })
 
- // Update the router with the history and context
- router.update({
- history,
- origin: router.options.origin ?? origin,
- })
+  // Update the router with the history and context
+  router.update({
+  history,
+  origin: router.options.origin ?? origin,
+  })
 
- await router.load()
+  await router.load()
 
- await router.serverSsr?.dehydrate()
+  await router.serverSsr?.dehydrate()
 
- const responseHeaders = getRequestHeaders({
- router,
- })
+  const responseHeaders = getRequestHeaders({
+  router,
+  })
 
- return cb({
- request,
- router,
- responseHeaders,
- } as any)
+ // Mark that the callback will handle cleanup
+ cbWillCleanup = true
+ return cb({
+ request,
+ router,
+ responseHeaders,
+ } as any)
+ } finally {
+ if (!cbWillCleanup) {
+ // Clean up router SSR state if the callback won't handle it
+ // (e.g., if an error occurred before the callback was invoked).
+ // When the callback runs, it handles cleanup (either via transformStreamWithRouter
+ // for streaming, or directly in renderRouterToString for non-streaming).
+ router.serverSsr?.cleanup()
+ }
+ }
  }
 }
 

packages/router-core/src/ssr/ssr-server.ts

@@ -16,6 +16,7 @@ import type { Manifest, RouterManagedTag } from '../manifest'
 declare module '../router' {
  interface ServerSsr {
  setRenderFinished: () => void
+ cleanup: () => void
  }
  interface RouterEvents {
  onInjectedHtml: {
@@ -55,11 +56,17 @@ const INITIAL_SCRIPTS = [
 ]
 
 class ScriptBuffer {
- constructor(private router: AnyRouter) {}
+ private router: AnyRouter | undefined
  private _queue: Array<string> = [...INITIAL_SCRIPTS]
  private _scriptBarrierLifted = false
+ private _cleanedUp = false
+
+ constructor(router: AnyRouter) {
+ this.router = router
+ }
 
  enqueue(script: string) {
+ if (this._cleanedUp) return
  if (this._scriptBarrierLifted && this._queue.length === 0) {
  queueMicrotask(() => {
  this.injectBufferedScripts()
@@ -69,7 +76,7 @@ class ScriptBuffer {
  }
 
  liftBarrier() {
- if (this._scriptBarrierLifted) return
+ if (this._scriptBarrierLifted || this._cleanedUp) return
  this._scriptBarrierLifted = true
  if (this._queue.length > 0) {
  queueMicrotask(() => {
@@ -90,11 +97,18 @@ class ScriptBuffer {
  }
 
  injectBufferedScripts() {
+ if (this._cleanedUp) return
  const scriptsToInject = this.takeAll()
- if (scriptsToInject) {
- this.router.serverSsr!.injectScript(() => scriptsToInject)
+ if (scriptsToInject && this.router?.serverSsr) {
+ this.router.serverSsr.injectScript(() => scriptsToInject)
  }
  }
+
+ cleanup() {
+ this._cleanedUp = true
+ this._queue = []
+ this.router = undefined
+ }
 }
 
 export function attachRouterServerSsrUtils({
@@ -203,6 +217,8 @@ export function attachRouterServerSsrUtils({
  onRenderFinished: (listener) => listeners.push(listener),
  setRenderFinished: () => {
  listeners.forEach((l) => l())
+ // Clear listeners after calling them to prevent memory leaks
+ listeners.length = 0
  scriptBuffer.liftBarrier()
  },
  takeBufferedScripts() {
@@ -221,6 +237,14 @@ export function attachRouterServerSsrUtils({
  liftScriptBarrier() {
  scriptBuffer.liftBarrier()
  },
+ cleanup() {
+ // Guard against multiple cleanup calls
+ if (!router.serverSsr) return
+ listeners.length = 0
+ scriptBuffer.cleanup()
+ router.serverSsr.injectedHtml = []
+ router.serverSsr = undefined
+ },
  }
 }
 

packages/router-core/src/ssr/transformStreamWithRouter.ts

@@ -35,7 +35,7 @@ type ReadablePassthrough = {
  destroyed: boolean
 }
 
-function createPassthrough(onCancel?: () => void) {
+function createPassthrough(onCancel: () => void) {
  let controller: ReadableStreamDefaultController<any>
  const encoder = new TextEncoder()
  const stream = new ReadableStream({
@@ -44,27 +44,33 @@ function createPassthrough(onCancel?: () => void) {
  },
  cancel() {
  res.destroyed = true
- onCancel?.()
+ onCancel()
  },
  })
 
  const res: ReadablePassthrough = {
  stream,
  write: (chunk) => {
+ // Don't write to destroyed stream
+ if (res.destroyed) return
  if (typeof chunk === 'string') {
  controller.enqueue(encoder.encode(chunk))
  } else {
  controller.enqueue(chunk)
  }
  },
  end: (chunk) => {
+ // Don't end already destroyed stream
+ if (res.destroyed) return
  if (chunk) {
  res.write(chunk)
  }
- controller.close()
  res.destroyed = true
+ controller.close()
  },
  destroy: (error) => {
+ // Don't destroy already destroyed stream
+ if (res.destroyed) return
  res.destroyed = true
  controller.error(error)
  },
@@ -105,14 +111,20 @@ export function transformStreamWithRouter(
 ) {
  let stopListeningToInjectedHtml: (() => void) | undefined = undefined
  let timeoutHandle: NodeJS.Timeout
+ let cleanedUp = false
 
- const finalPassThrough = createPassthrough(() => {
+ function cleanup() {
+ if (cleanedUp) return
+ cleanedUp = true
  if (stopListeningToInjectedHtml) {
  stopListeningToInjectedHtml()
  stopListeningToInjectedHtml = undefined
  }
  clearTimeout(timeoutHandle)
- })
+ router.serverSsr?.cleanup()
+ }
+
+ const finalPassThrough = createPassthrough(cleanup)
  const textDecoder = new TextDecoder()
 
  let isAppRendering = true
@@ -149,13 +161,16 @@ export function transformStreamWithRouter(
  )
 
  function handleInjectedHtml() {
+ // Don't process if already cleaned up
+ if (cleanedUp) return
+
  router.serverSsr!.injectedHtml.forEach((promise) => {
  processingCount++
 
  promise
  .then((html) => {
- // Don't write to destroyed stream
- if (finalPassThrough.destroyed) {
+ // Don't write to destroyed stream or after cleanup
+ if (cleanedUp || finalPassThrough.destroyed) {
  return
  }
  if (isAppRendering) {
@@ -165,10 +180,7 @@ export function transformStreamWithRouter(
  }
  })
  .catch((err) => {
- // Only reject if not already settled
- if (!finalPassThrough.destroyed) {
- injectedHtmlDonePromise.reject(err)
- }
+ injectedHtmlDonePromise.reject(err)
  })
  .finally(() => {
  processingCount--
@@ -183,6 +195,11 @@ export function transformStreamWithRouter(
 
  injectedHtmlDonePromise
  .then(() => {
+ // Don't process if already cleaned up or destroyed
+ if (cleanedUp || finalPassThrough.destroyed) {
+ return
+ }
+
  clearTimeout(timeoutHandle)
  const finalHtml =
  leftover + leftoverHtml + getBufferedRouterStream() + pendingClosingTags
@@ -194,19 +211,24 @@ export function transformStreamWithRouter(
  finalPassThrough.end(finalHtml)
  })
  .catch((err) => {
+ // Don't process if already cleaned up
+ if (cleanedUp || finalPassThrough.destroyed) {
+ return
+ }
+
  console.error('Error reading routerStream:', err)
  finalPassThrough.destroy(err)
  })
- .finally(() => {
- if (stopListeningToInjectedHtml) {
- stopListeningToInjectedHtml()
- stopListeningToInjectedHtml = undefined
- }
- })
+ .finally(cleanup)
 
  // Transform the appStream
  readStream(appStream, {
  onData: (chunk) => {
+ // Don't process if already cleaned up
+ if (cleanedUp || finalPassThrough.destroyed) {
+ return
+ }
+
  const text = decodeChunk(chunk.value)
  const chunkString = leftover + text
  const bodyEndMatch = chunkString.match(patternBodyEnd)
@@ -266,8 +288,8 @@ export function transformStreamWithRouter(
  }
  },
  onEnd: () => {
- // Don't process if stream was already destroyed/cancelled
- if (finalPassThrough.destroyed) {
+ // Don't process if stream was already destroyed/cancelled or cleaned up
+ if (cleanedUp || finalPassThrough.destroyed) {
  return
  }
 
@@ -288,6 +310,11 @@ export function transformStreamWithRouter(
  }
  },
  onError: (error) => {
+ // Don't process if already cleaned up
+ if (cleanedUp) {
+ return
+ }
+
  console.error('Error reading appStream:', error)
  isAppRendering = false
  router.serverSsr!.setRenderFinished()

packages/solid-router/src/ssr/renderRouterToString.tsx

@@ -40,5 +40,7 @@ export const renderRouterToString = async ({
  status: 500,
  headers: responseHeaders,
  })
+ } finally {
+ router.serverSsr?.cleanup()
  }
 }

packages/start-server-core/src/createStartHandler.ts

@@ -81,6 +81,8 @@ export function createStartHandler<TRegister = Register>(
  requestOpts,
  ) => {
  let router: AnyRouter | null = null as AnyRouter | null
+ // Track whether the callback will handle cleanup
+ let cbWillCleanup = false as boolean
  try {
  const origin = getOrigin(request)
 
@@ -207,6 +209,8 @@ export function createStartHandler<TRegister = Register>(
  await router.serverSsr!.dehydrate()
 
  const responseHeaders = getStartResponseHeaders({ router })
+ // Mark that the callback will handle cleanup
+ cbWillCleanup = true
  const response = await cb({
  request,
  router,
@@ -312,9 +316,14 @@ export function createStartHandler<TRegister = Register>(
 
  return response
  } finally {
- if (router) {
- router = null
+ if (router && !cbWillCleanup) {
+ // Clean up router SSR state if it was set up but won't be cleaned up by the callback
+ // (e.g., in redirect cases or early returns before the callback is invoked).
+ // When the callback runs, it handles cleanup (either via transformStreamWithRouter
+ // for streaming, or directly in renderRouterToString for non-streaming).
+ router.serverSsr?.cleanup()
  }
+ router = null
  }
  }