fix(mobile): skip queued generation stages after cancellation (#34809)

* fix(mobile): skip queued generation stages after cancellation

* test(mobile): fold canceled queued-stage case into download suite

---------

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

mobile/src/mobile-web-shell/generation-store.ts

@@ -42,7 +42,11 @@ export type StagedGeneration = {
 export type GenerationStore = {
   readActiveGeneration(hostKey: string): Promise<ActiveGeneration | null>
   openReusableAssets(hostKey: string): Promise<MobileWebBundleReusableAssets>
-  stageGeneration(hostKey: string, result: MobileWebBundleFetchResult): Promise<StagedGeneration>
+  stageGeneration(
+    hostKey: string,
+    result: MobileWebBundleFetchResult,
+    signal?: AbortSignal
+  ): Promise<StagedGeneration>
   commitGeneration(staged: StagedGeneration): Promise<ActiveGeneration>
   abortStagedGeneration(staged: StagedGeneration): Promise<void>
   sweepStagedGenerations(): Promise<void>
@@ -157,8 +161,12 @@ export function createGenerationStore(options: {
 
   async function stage(
     hostKey: string,
-    result: MobileWebBundleFetchResult
+    result: MobileWebBundleFetchResult,
+    signal?: AbortSignal
   ): Promise<StagedGeneration> {
+    if (signal?.aborted) {
+      throw new Error('generation staging was aborted')
+    }
     const manifest = result.manifest
     const directory = joinUri(stagingRoot(hostKey), manifest.buildId)
     const reused = new Set(result.reused)
@@ -272,7 +280,7 @@ export function createGenerationStore(options: {
   return {
     readActiveGeneration: (hostKey) => serialize(() => openActive(hostKey)),
     openReusableAssets: (hostKey) => serialize(() => openReusable(hostKey)),
-    stageGeneration: (hostKey, result) => serialize(() => stage(hostKey, result)),
+    stageGeneration: (hostKey, result, signal) => serialize(() => stage(hostKey, result, signal)),
     commitGeneration: (staged) => serialize(() => commit(staged)),
     abortStagedGeneration: (staged) =>
       serialize(() => fs.delete(requireIssuedHandle(staged).directory)),

mobile/src/mobile-web-shell/mobile-web-shell-session-effects.test.ts

@@ -255,7 +255,7 @@ describe('a download that reuses the cached generation', () => {
     return { fileSystem, store }
   }
 
-  function runFor(store: GenerationStore) {
+  function runFor(store: GenerationStore, downloads = new Set<AbortController>()) {
     const { events, send } = collect()
     const done = download({
       client: CLIENT,
@@ -269,12 +269,70 @@ describe('a download that reuses the cached generation', () => {
         setTimer: () => () => {}
       },
       startedAt: 0,
-      downloads: new Set(),
+      downloads,
       send
     })
     return done.then(() => events)
   }
 
+  it('writes no staging bytes when canceled while its staging waits in the store queue', async () => {
+    const fileSystem = createFakeGenerationFileSystem()
+    const store = createGenerationStore({ fileSystem })
+    const NEXT = bundle(16)
+    const BLOCKER = deriveHostCacheKey('host-b')
+    let enterBlocker = (): void => {}
+    const blockerEntered = new Promise<void>((resolve) => {
+      enterBlocker = resolve
+    })
+    let releaseBlocker = (): void => {}
+    const blockerReleased = new Promise<void>((resolve) => {
+      releaseBlocker = resolve
+    })
+    const list = fileSystem.list.bind(fileSystem)
+    fileSystem.list = async (uri) => {
+      if (uri.includes(`/${BLOCKER}/`)) {
+        enterBlocker()
+        await blockerReleased
+      }
+      return list(uri)
+    }
+    const writes: string[] = []
+    const writeBytes = fileSystem.writeBytes.bind(fileSystem)
+    fileSystem.writeBytes = async (uri, bytes) => {
+      writes.push(uri)
+      await writeBytes(uri, bytes)
+    }
+    let blocker: Promise<unknown> = Promise.resolve()
+    doubles.fetch = async () => {
+      blocker = store.readActiveGeneration(BLOCKER)
+      await blockerEntered
+      return NEXT
+    }
+    let markQueued = (): void => {}
+    const stageQueued = new Promise<void>((resolve) => {
+      markQueued = resolve
+    })
+    const downloads = new Set<AbortController>()
+    const run = runFor(
+      {
+        ...store,
+        stageGeneration: (...args) => {
+          const staged = store.stageGeneration(...args)
+          markQueued()
+          return staged
+        }
+      },
+      downloads
+    )
+
+    await stageQueued
+    downloads.forEach((controller) => controller.abort())
+    releaseBlocker()
+    await Promise.all([run, blocker])
+
+    expect(writes.filter((uri) => uri.includes(`/tmp/${NEXT.manifest.buildId}/`))).toEqual([])
+  })
+
   it('commits whole when another mount sweeps the staging trees mid-download', async () => {
     const { fileSystem, store } = await cachedStore()
     doubles.fetch = async () => {

mobile/src/mobile-web-shell/mobile-web-shell-session-effects.ts

@@ -114,7 +114,7 @@ export async function download(args: {
     }
     fetchedWhole = true
     send({ type: 'download-staged', flow })
-    const staged = await store.stageGeneration(hostKey, fetched)
+    const staged = await store.stageGeneration(hostKey, fetched, controller.signal)
     // Recheck after staging: abort may remove tmp, but cannot undo a commit’s activation/index changes.
     if (controller.signal.aborted) {
       await store.abortStagedGeneration(staged).catch(() => undefined)