Stop obsolete worktree preparations when evicted or expired

src/main/git/worktree-create-preparation-real-git.test.ts

@@ -1,8 +1,10 @@
 import { execFileSync } from 'node:child_process'
+import { existsSync, watch, type FSWatcher } from 'node:fs'
 import { mkdir, mkdtemp, readFile, realpath, rm, writeFile } from 'node:fs/promises'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
-import { afterEach, describe, expect, it } from 'vitest'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import * as gitRunner from './runner'
 import {
   createWorktreePreparationLockReason,
   isWorktreeCreatePreparation,
@@ -46,6 +48,59 @@ afterEach(async () => {
 })
 
 describe('prepared worktree creation with real Git', () => {
+  it('removes partial checkout files and registration after materialization is aborted', async () => {
+    const { repoPath, root } = await createRepo()
+    await Promise.all(
+      Array.from({ length: 1000 }, (_, index) =>
+        writeFile(
+          join(repoPath, `payload-${index.toString().padStart(4, '0')}.txt`),
+          'payload'.repeat(128)
+        )
+      )
+    )
+    git(repoPath, ['add', '.'])
+    git(repoPath, ['commit', '--quiet', '-m', 'materialization fixture'])
+    const preparationRoot = join(root, WORKTREE_CREATE_PREPARATION_DIRECTORY)
+    const preparedPath = join(preparationRoot, `${process.pid}-partial`)
+    await mkdir(preparationRoot, { recursive: true })
+    const controller = new AbortController()
+    const original = gitRunner.gitExecFileAsync
+    let watcher: FSWatcher | undefined
+    let observedMaterialization = false
+    const calls: string[][] = []
+    const spy = vi.spyOn(gitRunner, 'gitExecFileAsync').mockImplementation((args, options) => {
+      calls.push([...args])
+      if (args.includes('reset')) {
+        watcher = watch(preparedPath, (_event, filename) => {
+          if (filename?.toString().startsWith('payload-')) {
+            observedMaterialization = true
+            watcher?.close()
+            controller.abort()
+          }
+        })
+      }
+      return original(args, options)
+    })
+    try {
+      await expect(
+        prepareWorktreeCreateCheckout(
+          repoPath,
+          preparedPath,
+          'main',
+          createWorktreePreparationLockReason('partial-test'),
+          { signal: controller.signal }
+        )
+      ).rejects.toThrow()
+      expect(observedMaterialization).toBe(true)
+      expect(calls.some((args) => args[args.indexOf('worktree') + 1] === 'lock')).toBe(false)
+      expect(existsSync(preparedPath)).toBe(false)
+      expect(await listWorktrees(repoPath, { includeCreatePreparations: true })).toHaveLength(1)
+    } finally {
+      watcher?.close()
+      spy.mockRestore()
+    }
+  })
+
   it('cleans up when the create signal is canceled', async () => {
     const { repoPath, root } = await createRepo()
     const preparationRoot = join(root, WORKTREE_CREATE_PREPARATION_DIRECTORY)

src/main/worktree-create-preparation-pool.ts

@@ -38,6 +38,8 @@ export type PreparationEntry = {
   createdAt: number
   ready: Promise<void>
   expiration: NodeJS.Timeout
+  controller: AbortController
+  checkoutStarted: boolean
 }
 
 export type StartPreparationArgs = {
@@ -68,6 +70,9 @@ async function discardEntry(entry: PreparationEntry): Promise<void> {
   // A failed checkout self-discards, but that self-discard is best-effort too, so it can strand the
   // registration for the same reason the discard here can. Enrol either way.
   await entry.ready.catch(() => {})
+  if (!entry.checkoutStarted) {
+    return
+  }
   await discardPreparationWithRetry({
     hostKey: preparationHostKey(entry.repoPathKey, entry.wslDistro),
     repoPath: entry.repoPath,
@@ -86,6 +91,7 @@ function expireEntry(entry: PreparationEntry): void {
     return
   }
   preparations.delete(entry.key)
+  entry.controller.abort()
   discardEntryInBackground(entry)
 }
 
@@ -118,6 +124,7 @@ function enforcePreparationLimit(
     }
     preparations.delete(victim.key)
     clearTimeout(victim.expiration)
+    victim.controller.abort()
     discardEntryInBackground(victim)
   }
 }
@@ -163,6 +170,10 @@ export function startPreparation({
     WORKTREE_CREATE_PREPARATION_DIRECTORY
   )
   const preparedPath = pathOps(workspaceRoot).join(preparationRoot, preparationId)
+  const controller = new AbortController()
+  const signal = options.signal
+    ? AbortSignal.any([options.signal, controller.signal])
+    : controller.signal
   const entry = {} as PreparationEntry
   const expiration = setTimeout(() => expireEntry(entry), WORKTREE_CREATE_PREPARATION_TTL_MS)
   expiration.unref()
@@ -179,17 +190,19 @@ export function startPreparation({
     options,
     createdAt: Date.now(),
     expiration,
+    controller,
+    checkoutStarted: false,
     ready: (async () => {
       await cleanupStalePreparations(preparationHostKey(repoPathKey, wslDistro), repoPath, options)
+      signal.throwIfAborted()
       await mkdir(toHostFilesystemPath(preparationRoot), { recursive: true })
+      signal.throwIfAborted()
       // Already canonical, so the add re-resolves nothing.
-      await prepareWorktreeCreateCheckout(
-        repoPath,
-        preparedPath,
-        canonicalBase,
-        lockReason,
-        options
-      )
+      entry.checkoutStarted = true
+      await prepareWorktreeCreateCheckout(repoPath, preparedPath, canonicalBase, lockReason, {
+        ...options,
+        signal
+      })
     })()
   } satisfies PreparationEntry)
   preparations.set(key, entry)

src/main/worktree-create-preparation.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import type { Store } from './persistence'
+import { WORKTREE_CREATE_PREPARATION_TTL_MS } from './worktree-create-preparation-pool'
 import type { Repo } from '../shared/repo-types'
 import { WORKTREE_CREATE_PREPARATION_DIRECTORY } from '../shared/worktree/create-preparation'
 import { resolveWorktreeAddBaseRef } from '../shared/worktree/base-ref'
@@ -96,6 +97,124 @@ afterEach(async () => {
 })
 
 describe('worktree create preparation registry', () => {
+  it('cancels an evicted checkout and cleans up with the original options', async () => {
+    let signal: AbortSignal | undefined
+    mocks.prepareCheckout.mockImplementationOnce((_repo, _path, _base, _lock, options) => {
+      signal = options.signal
+      return new Promise<void>((_resolve, reject) => {
+        signal!.addEventListener('abort', () => reject(signal!.reason), { once: true })
+      })
+    })
+    const obsolete = prepareWorktreeCreateForRepo(store, repo, 'origin/main')
+    const settled = Promise.allSettled([obsolete])
+    await flushBackgroundWork()
+    const obsoletePath = mocks.prepareCheckout.mock.calls[0][1]
+    for (const base of ['origin/one', 'origin/two', 'origin/three']) {
+      await prepareWorktreeCreateForRepo(store, repo, base)
+    }
+    expect(signal?.aborted).toBe(true)
+    expect((await settled)[0].status).toBe('rejected')
+    await flushBackgroundWork()
+    expect(mocks.discard).toHaveBeenCalledWith(repo.path, obsoletePath, {})
+  })
+
+  it('does not start obsolete checkout work after shared cleanup finishes', async () => {
+    let releaseCleanup!: () => void
+    mocks.listWorktreeGraph.mockImplementationOnce(
+      () =>
+        new Promise<[]>((resolve) => {
+          releaseCleanup = () => resolve([])
+        })
+    )
+    const requests = ['main', 'one', 'two', 'three'].map((base) =>
+      prepareWorktreeCreateForRepo(store, repo, `origin/${base}`)
+    )
+    const settled = Promise.allSettled(requests)
+    await flushBackgroundWork()
+    expect(mocks.prepareCheckout).not.toHaveBeenCalled()
+    releaseCleanup()
+    const results = await settled
+    expect(results.map((result) => result.status)).toEqual([
+      'rejected',
+      'fulfilled',
+      'fulfilled',
+      'fulfilled'
+    ])
+    expect(mocks.prepareCheckout).toHaveBeenCalledTimes(3)
+    await flushBackgroundWork()
+    expect(mocks.discard).not.toHaveBeenCalled()
+  })
+
+  it('keeps a claimed in-flight checkout alive when new preparations fill the pool', async () => {
+    let signal: AbortSignal | undefined
+    let finishCheckout!: () => void
+    mocks.prepareCheckout.mockImplementationOnce((_repo, _path, _base, _lock, options) => {
+      signal = options.signal
+      return new Promise<void>((resolve) => {
+        finishCheckout = resolve
+      })
+    })
+    const preparation = prepareWorktreeCreateForRepo(store, repo, 'origin/main')
+    await flushBackgroundWork()
+    const create = consumePreparedWorktreeCreate({
+      repoPath: repo.path,
+      workspaceRoot: '/workspace',
+      worktreePath: '/workspace/claimed',
+      branch: 'claimed',
+      baseBranch: 'origin/main'
+    })
+    await flushBackgroundWork()
+    for (const base of ['origin/one', 'origin/two', 'origin/three', 'origin/four']) {
+      await prepareWorktreeCreateForRepo(store, repo, base)
+    }
+    expect(signal?.aborted).toBe(false)
+    finishCheckout()
+    await preparation
+    expect(await create).toMatchObject({ status: 'hit' })
+  })
+
+  it('cancels an expired in-flight checkout', async () => {
+    vi.useFakeTimers()
+    let signal: AbortSignal | undefined
+    mocks.prepareCheckout.mockImplementationOnce((_repo, _path, _base, _lock, options) => {
+      signal = options.signal
+      return new Promise<void>((_resolve, reject) => {
+        signal!.addEventListener('abort', () => reject(signal!.reason), { once: true })
+      })
+    })
+    try {
+      const settled = Promise.allSettled([prepareWorktreeCreateForRepo(store, repo, 'origin/main')])
+      await vi.advanceTimersByTimeAsync(0)
+      expect(signal?.aborted).toBe(false)
+      await vi.advanceTimersByTimeAsync(WORKTREE_CREATE_PREPARATION_TTL_MS)
+      expect(signal?.aborted).toBe(true)
+      expect((await settled)[0].status).toBe('rejected')
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('preserves caller cancellation without mutating its options', async () => {
+    const controller = new AbortController()
+    const options = { signal: controller.signal }
+    mocks.getWorktreeOptions.mockReturnValue(options)
+    let signal: AbortSignal | undefined
+    mocks.prepareCheckout.mockImplementationOnce((_repo, _path, _base, _lock, executionOptions) => {
+      signal = executionOptions.signal
+      return new Promise<void>((_resolve, reject) => {
+        signal!.addEventListener('abort', () => reject(signal!.reason), { once: true })
+      })
+    })
+    const preparation = prepareWorktreeCreateForRepo(store, repo, 'origin/main')
+    const settled = Promise.allSettled([preparation])
+    await flushBackgroundWork()
+    controller.abort()
+    expect(signal?.aborted).toBe(true)
+    expect((await settled)[0].status).toBe('rejected')
+    expect(options.signal).toBe(controller.signal)
+    expect(signal).not.toBe(controller.signal)
+  })
+
   it('starts the checkout only once the async workspace root resolves', async () => {
     let resolveRoot!: (root: string) => void
     mocks.computeWorkspaceRootAsync.mockReturnValue(
@@ -364,7 +483,7 @@ describe('worktree create preparation registry', () => {
       expect.any(String),
       'refs/remotes/origin/main',
       expect.any(String),
-      options
+      { ...options, signal: expect.any(AbortSignal) }
     )
     expect(mocks.finalize).toHaveBeenCalledWith(
       repo.path,