Serve stale docs artifacts while refreshing

Author: tannerlinsley

src/server.ts

@@ -4,7 +4,10 @@ import { wrapFetchWithSentry } from '@sentry/tanstackstart-react'
 import handler, { createServerEntry } from '@tanstack/react-start/server-entry'
 import { runWithDatabaseContext } from '~/db/client'
 import { runScheduledTasks } from '~/server/scheduled.server'
-import { runWithHostRuntimeEnv } from '~/server/runtime/host.server'
+import {
+  runWithHostRuntimeContext,
+  runWithHostRuntimeEnv,
+} from '~/server/runtime/host.server'
 import {
   installProductionFetchProbe,
   installProductionProcessProbe,
@@ -189,8 +192,10 @@ const server = createServerEntry(
 )
 
 export default {
-  fetch(request: Request, env: unknown) {
-    return runWithHostRuntimeEnv(env, () => server.fetch(request))
+  fetch(request: Request, env: unknown, context: unknown) {
+    return runWithHostRuntimeEnv(env, () =>
+      runWithHostRuntimeContext(context, () => server.fetch(request)),
+    )
   },
   scheduled(
     controller: ScheduledController,
@@ -199,8 +204,10 @@ export default {
   ) {
     context.waitUntil(
       runWithHostRuntimeEnv(env, () =>
-        runWithDatabaseContext(() =>
-          runScheduledTasks(controller.cron, controller.scheduledTime),
+        runWithHostRuntimeContext(context, () =>
+          runWithDatabaseContext(() =>
+            runScheduledTasks(controller.cron, controller.scheduledTime),
+          ),
         ),
       ),
     )

src/server/runtime/host.server.ts

@@ -19,7 +19,12 @@ type HostRuntimeModule = {
   env: HostRuntimeEnv
 }
 
+type HostRuntimeContext = {
+  waitUntil(promise: Promise<unknown>): void
+}
+
 const hostRuntimeEnvStorage = new AsyncLocalStorage<HostRuntimeEnv>()
+const hostRuntimeContextStorage = new AsyncLocalStorage<HostRuntimeContext>()
 
 function isObject(value: unknown): value is Record<string, unknown> {
   return typeof value === 'object' && value !== null
@@ -29,6 +34,14 @@ function isHostRuntimeEnv(value: unknown): value is HostRuntimeEnv {
   return isObject(value)
 }
 
+function isHostRuntimeContext(value: unknown): value is HostRuntimeContext {
+  return (
+    isObject(value) &&
+    'waitUntil' in value &&
+    typeof value.waitUntil === 'function'
+  )
+}
+
 function isStaticAssetService(value: unknown): value is StaticAssetService {
   return (
     isObject(value) && 'fetch' in value && typeof value.fetch === 'function'
@@ -62,10 +75,29 @@ export function runWithHostRuntimeEnv<T>(env: unknown, fn: () => T): T {
   return hostRuntimeEnvStorage.run(env, fn)
 }
 
+export function runWithHostRuntimeContext<T>(context: unknown, fn: () => T): T {
+  if (!isHostRuntimeContext(context)) {
+    return fn()
+  }
+
+  return hostRuntimeContextStorage.run(context, fn)
+}
+
 export function getCurrentHostRuntimeEnv() {
   return hostRuntimeEnvStorage.getStore()
 }
 
+export function scheduleHostRuntimeTask(createTask: () => Promise<unknown>) {
+  const context = hostRuntimeContextStorage.getStore()
+
+  if (!context) {
+    return false
+  }
+
+  context.waitUntil(createTask())
+  return true
+}
+
 async function getStaticAssetService() {
   if (!isIsolateRuntime()) return
 

src/utils/github-content-cache.server.ts

@@ -5,6 +5,7 @@ import {
   type BlobStorageCache,
   type BlobStorageListedObject,
 } from '~/server/runtime/blob-storage.server'
+import { scheduleHostRuntimeTask } from '~/server/runtime/host.server'
 import { isValidRepoPath, MAX_REPO_PATH_LENGTH } from './repo-path'
 
 const POSITIVE_STALE_MS = 5 * 60 * 1000
@@ -256,10 +257,10 @@ function isFresh(staleAt: Date) {
 
 // markGitHubContentStale / markDocsArtifactsStale set staleAt to the epoch
 // (new Date(0)) as a sentinel for "forcibly invalidated": an admin clicked
-// the purge button or a push webhook fired. Natural TTL expiry and forced
-// invalidation both refresh synchronously now. The object stays around so the
-// bottom of getCachedGitHubContent / getCachedDocsArtifact can still fall back
-// to it if GitHub is unreachable.
+// the purge button or a push webhook fired. The object stays around so refresh
+// paths can still fall back to it if GitHub is unreachable, and docs artifacts
+// can serve it while scheduling a background rebuild when the runtime supports
+// waitUntil.
 function isForciblyStale(staleAt: Date) {
   return staleAt.getTime() <= 0
 }
@@ -1408,43 +1409,65 @@ export async function getCachedDocsArtifact<T>(opts: {
     cachedRow && opts.isValue(cachedRow.payload) ? cachedRow.payload : undefined
   const forciblyStale = !!cachedRow && isForciblyStale(cachedRow.staleAt)
 
-  if (storedValue !== undefined && !forciblyStale) {
-    if (cachedRow && isFresh(cachedRow.staleAt)) {
+  const refreshArtifact = () =>
+    withPendingRefresh(cacheKey, async () => {
+      const latestRow = await readRow()
+      const latestValue =
+        latestRow && opts.isValue(latestRow.payload)
+          ? latestRow.payload
+          : undefined
+      const latestForciblyStale =
+        !!latestRow && isForciblyStale(latestRow.staleAt)
+
+      if (
+        latestValue !== undefined &&
+        latestRow &&
+        !latestForciblyStale &&
+        isFresh(latestRow.staleAt)
+      ) {
+        return latestValue
+      }
+
+      try {
+        const payload = await opts.build()
+        await upsertDocsArtifact({ ...opts, payload })
+        return payload
+      } catch (error) {
+        if (latestValue !== undefined) {
+          console.warn(`[GitHub Cache] Serving stale artifact ${cacheKey}`)
+          return latestValue
+        }
+
+        throw error
+      }
+    })
+
+  if (storedValue !== undefined) {
+    if (!forciblyStale && cachedRow && isFresh(cachedRow.staleAt)) {
       return storedValue
     }
-  }
-
-  return withPendingRefresh(cacheKey, async () => {
-    const latestRow = await readRow()
-    const latestValue =
-      latestRow && opts.isValue(latestRow.payload)
-        ? latestRow.payload
-        : undefined
-    const latestForciblyStale =
-      !!latestRow && isForciblyStale(latestRow.staleAt)
 
     if (
-      latestValue !== undefined &&
-      latestRow &&
-      !latestForciblyStale &&
-      isFresh(latestRow.staleAt)
+      scheduleHostRuntimeTask(() =>
+        refreshArtifact().then(
+          () => undefined,
+          (error) => {
+            console.warn(
+              `[GitHub Cache] Background artifact refresh failed ${cacheKey}`,
+              error,
+            )
+          },
+        ),
+      )
     ) {
-      return latestValue
+      console.warn(
+        `[GitHub Cache] Serving stale artifact ${cacheKey}; refreshing in background`,
+      )
+      return storedValue
     }
+  }
 
-    try {
-      const payload = await opts.build()
-      await upsertDocsArtifact({ ...opts, payload })
-      return payload
-    } catch (error) {
-      if (latestValue !== undefined) {
-        console.warn(`[GitHub Cache] Serving stale artifact ${cacheKey}`)
-        return latestValue
-      }
-
-      throw error
-    }
-  })
+  return refreshArtifact()
 }
 
 export async function listDocsCacheRepoStats() {

tests/github-content-cache.test.ts

@@ -1,5 +1,8 @@
 import assert from 'node:assert/strict'
-import { runWithHostRuntimeEnv } from '../src/server/runtime/host.server'
+import {
+  runWithHostRuntimeContext,
+  runWithHostRuntimeEnv,
+} from '../src/server/runtime/host.server'
 import {
   getCachedDocsArtifact,
   getCachedGitHubJsonContent,
@@ -22,6 +25,15 @@ function isObject(value: unknown): value is Record<string, unknown> {
   return typeof value === 'object' && value !== null
 }
 
+async function withTimeout(promise: Promise<void>, message: string) {
+  await Promise.race([
+    promise,
+    new Promise<void>((_, reject) => {
+      setTimeout(() => reject(new Error(message)), 1000)
+    }),
+  ])
+}
+
 function isDocsManifest(value: unknown): value is { paths: Array<string> } {
   if (!isObject(value)) {
     return false
@@ -294,6 +306,85 @@ async function testArtifactInvalidationAndPruneDelete() {
   assert.equal(originCalls, 1)
 }
 
+async function testStaleArtifactRefreshesInWaitUntil() {
+  resetGitHubContentCacheForTest()
+
+  await getCachedDocsArtifact({
+    repo,
+    gitRef,
+    docsRoot: 'docs',
+    artifactType: 'docs-manifest',
+    artifactKey: 'default',
+    isValue: isDocsManifest,
+    build: async () => ({ paths: ['old'] }),
+  })
+
+  assert.equal(await markDocsArtifactsStale({ repo, gitRef }), 1)
+
+  const waitUntilPromises: Array<Promise<unknown>> = []
+  let buildCalls = 0
+  let resolveRefresh: ((value: { paths: Array<string> }) => void) | undefined
+  let markRefreshStarted: (() => void) | undefined
+  const refreshStarted = new Promise<void>((resolve) => {
+    markRefreshStarted = resolve
+  })
+  const refreshResult = new Promise<{ paths: Array<string> }>((resolve) => {
+    resolveRefresh = resolve
+  })
+
+  const staleResult = await runWithHostRuntimeContext(
+    {
+      waitUntil(promise: Promise<unknown>) {
+        waitUntilPromises.push(promise)
+      },
+    },
+    () =>
+      getCachedDocsArtifact({
+        repo,
+        gitRef,
+        docsRoot: 'docs',
+        artifactType: 'docs-manifest',
+        artifactKey: 'default',
+        isValue: isDocsManifest,
+        build: () => {
+          buildCalls += 1
+          markRefreshStarted?.()
+          return refreshResult
+        },
+      }),
+  )
+
+  assert.deepEqual(staleResult, { paths: ['old'] })
+  assert.equal(waitUntilPromises.length, 1)
+
+  await withTimeout(
+    refreshStarted,
+    'stale artifact refresh did not start in waitUntil',
+  )
+  assert.equal(buildCalls, 1)
+
+  if (!resolveRefresh) {
+    throw new Error('stale artifact refresh resolver was not created')
+  }
+
+  resolveRefresh({ paths: ['new'] })
+  await Promise.all(waitUntilPromises)
+
+  const freshResult = await getCachedDocsArtifact({
+    repo,
+    gitRef,
+    docsRoot: 'docs',
+    artifactType: 'docs-manifest',
+    artifactKey: 'default',
+    isValue: isDocsManifest,
+    build: async () => {
+      throw new Error('fresh artifact should not rebuild')
+    },
+  })
+
+  assert.deepEqual(freshResult, { paths: ['new'] })
+}
+
 await testMissStoresContent()
 await testFreshHitSkipsOrigin()
 await testWorkerEnvUsesBlobStorage()
@@ -304,5 +395,6 @@ await testJsonNegativeEntryRefreshes()
 await testRefreshFailureFallsBackToStaleContent()
 await testJsonContentUsesTypeGuard()
 await testArtifactInvalidationAndPruneDelete()
+await testStaleArtifactRefreshesInWaitUntil()
 
 console.log('github-content-cache tests passed')