From deab74e7ff109eb0f43571c3c5f12467733bf1c1 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:47:44 -0700 Subject: [PATCH] Stop obsolete worktree preparations when evicted or expired --- ...rktree-create-preparation-real-git.test.ts | 57 ++++++++- src/main/worktree-create-preparation-pool.ts | 27 +++- src/main/worktree-create-preparation.test.ts | 121 +++++++++++++++++- 3 files changed, 196 insertions(+), 9 deletions(-) diff --git a/src/main/git/worktree-create-preparation-real-git.test.ts b/src/main/git/worktree-create-preparation-real-git.test.ts index 63f8021a7ae..c593e1a57f4 100644 --- a/src/main/git/worktree-create-preparation-real-git.test.ts +++ b/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) diff --git a/src/main/worktree-create-preparation-pool.ts b/src/main/worktree-create-preparation-pool.ts index 7539c6076e4..1581f404b20 100644 --- a/src/main/worktree-create-preparation-pool.ts +++ b/src/main/worktree-create-preparation-pool.ts @@ -38,6 +38,8 @@ export type PreparationEntry = { createdAt: number ready: Promise expiration: NodeJS.Timeout + controller: AbortController + checkoutStarted: boolean } export type StartPreparationArgs = { @@ -68,6 +70,9 @@ async function discardEntry(entry: PreparationEntry): Promise { // 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) diff --git a/src/main/worktree-create-preparation.test.ts b/src/main/worktree-create-preparation.test.ts index 06818fec422..eec865d6936 100644 --- a/src/main/worktree-create-preparation.test.ts +++ b/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((_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((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((_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((_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,