Improve Cleanup of VS Code Sessions

vscode-dotnet-runtime-extension/src/extension.ts

@@ -136,6 +136,7 @@ const defaultTimeoutValue = 600;
 const moreInfoUrl = 'https://github.com/dotnet/vscode-dotnet-runtime/blob/main/Documentation/troubleshooting-runtime.md';
 let disableActivationUnderTest = true;
 let extensionEventStream: IEventStream | undefined;
+let extensionGlobalState: vscode.Memento | undefined;
 
 export function activate(vsCodeContext: vscode.ExtensionContext, extensionContext?: IExtensionContext)
 {
@@ -1001,6 +1002,10 @@ Installation will timeout in ${timeoutValue} seconds.`))
     // Preemptively install .NET for extensions who tell us to in their package.json
     const jsonInstaller = new JsonInstaller(globalEventStream, vsCodeExtensionContext);
 
+    // Store references for deactivate()
+    extensionEventStream = globalEventStream;
+    extensionGlobalState = vsCodeContext.globalState;
+
     // Exposing API Endpoints
     vsCodeContext.subscriptions.push(
         dotnetAcquireRegistration,
@@ -1031,3 +1036,13 @@ export function ReEnableActivationForManualActivation()
 {
     disableActivationUnderTest = false;
 }
+
+export async function deactivate(): Promise<void>
+{
+    if (extensionEventStream && extensionGlobalState)
+    {
+        const tracker = InstallTrackerSingleton.getInstance(extensionEventStream, extensionGlobalState);
+        await tracker.pruneStaleSessions();
+        await tracker.removeCurrentSession();
+    }
+}

vscode-dotnet-runtime-extension/src/test/functional/DotnetCoreAcquisitionExtension.test.ts

@@ -926,4 +926,38 @@ Paths: 'acquire returned: ${resultForAcquiringPathSettingRuntime.dotnetPath} whi
         assert.exists(result2!.dotnetPath, 'Second install should return a path');
         assert.include(result2!.dotnetPath, '9.0.0', 'Path should include the fully specified version 9.0.0');
     }).timeout(standardTimeoutTime);
+
+    test('Deactivate cleans up current session and prunes stale sessions from state', async () =>
+    {
+        const sessionStateKey = 'dotnet.returnedInstallDirectories';
+
+        // Step 1: Install a runtime so the current session is tracked in state
+        await installRuntime('9.0', 'runtime');
+
+        // Verify the current session was registered
+        const stateAfterInstall = mockState.get<Record<string, string[]>>(sessionStateKey, {});
+        const sessionIds = Object.keys(stateAfterInstall);
+        assert.isAtLeast(sessionIds.length, 1, 'At least one session should be tracked after install');
+
+        // Step 2: Seed a fake dead session (no process holds its mutex)
+        const fakeDeadSessionId = 'session-deadbeef00';
+        stateAfterInstall[fakeDeadSessionId] = ['/fake/path/dotnet'];
+        await mockState.update(sessionStateKey, stateAfterInstall);
+
+        const stateWithFake = mockState.get<Record<string, string[]>>(sessionStateKey, {});
+        assert.property(stateWithFake, fakeDeadSessionId, 'Fake dead session should be present in state before deactivate');
+
+        // Step 3: Call deactivate — should prune the dead session AND remove the current session
+        await extension.deactivate();
+
+        // Step 4: Verify the state is cleaned up
+        const stateAfterDeactivate = mockState.get<Record<string, string[]>>(sessionStateKey, {});
+        assert.notProperty(stateAfterDeactivate, fakeDeadSessionId, 'Fake dead session should have been pruned by deactivate');
+
+        // The current session should also be removed
+        for (const id of sessionIds)
+        {
+            assert.notProperty(stateAfterDeactivate, id, `Current session ${id} should have been removed by deactivate`);
+        }
+    }).timeout(standardTimeoutTime);
 });

vscode-dotnet-runtime-extension/src/test/functional/index.ts

@@ -6,7 +6,8 @@ import * as glob from 'glob';
 import * as Mocha from 'mocha';
 import * as path from 'path';
 
-export function run(): Promise<void> {
+export function run(): Promise<void>
+{
   // Create the mocha test
   const mocha = new Mocha({
     ui: 'tdd',
@@ -15,25 +16,33 @@ export function run(): Promise<void> {
 
   const testsRoot = path.resolve(__dirname, '..');
 
-  return new Promise((c, e) => {
-    glob('**/functional/**.test.js', { cwd: testsRoot }, (err, files) => {
-      if (err) {
+  return new Promise((c, e) =>
+  {
+    glob('**/functional/**.test.js', { cwd: testsRoot }, (err, files) =>
+    {
+      if (err)
+      {
         return e(err);
       }
 
       // Add files to the test suite
       files.forEach(f => mocha.addFile(path.resolve(testsRoot, f)));
 
-      try {
+      try
+      {
         // Run the mocha test
-        mocha.run(failures => {
-          if (failures > 0) {
+        mocha.run(failures =>
+        {
+          if (failures > 0)
+          {
             e(new Error(`${failures} tests failed.`));
-          } else {
+          } else
+          {
             c();
           }
         });
-      } catch (err) {
+      } catch (err)
+      {
         e(err);
       }
     });

vscode-dotnet-runtime-library/src/Acquisition/InstallTrackerSingleton.ts

@@ -16,6 +16,8 @@ import
     LiveDependentInUse,
     MarkedInstallInUse,
     ProcessEnvironmentCheck,
+    PrunedStaleSessions,
+    RemovingCurrentSession,
     RemovingExtensionFromList,
     RemovingOwnerFromList,
     RemovingVersionFromExtensionState,
@@ -198,22 +200,16 @@ export class InstallTrackerSingleton
                         return false; // Our session must be live if this code is running.
                     }
 
-                    // See if the session is still 'live' - there is no way to ensure we remove it on exit/os crash
-                    const logger = new EventStreamNodeIPCMutexLoggerWrapper(this.eventStream, sessionId);
-                    const mutex = new NodeIPCMutex(sessionId, logger, ``);
-
-                    const shouldContinue = await mutex.acquire(async () =>
+                    const isDead = await this.isSessionDead(sessionId);
+                    if (isDead)
                     {
-                        // eslint-disable-next-line no-return-await
-                        this.eventStream.post(new DependentIsDead(`Dependent Session ${sessionId} is no longer live - continue searching dependents.`))
                         existingSessionsWithUsedExecutablePaths.delete(sessionId);
                         await this.extensionState.update(this.sessionInstallsKey, serializeMapOfSets(existingSessionsWithUsedExecutablePaths));
-                        return Promise.resolve(true);
-                    }, 10, 30, `${sessionId}-${crypto.randomUUID()}`).catch(() => { return false; });
-                    if (!shouldContinue && exePaths.has(installExePath))
+                    }
+                    else if (exePaths.has(installExePath))
                     {
                         this.eventStream.post(new LiveDependentInUse(`Install ${installExePath} is in use by session ${sessionId}, so we can't uninstall it.`))
-                        return false; // We couldn't acquire the mutex, so the session must be live
+                        return false;
                     }
                 }
 
@@ -464,6 +460,80 @@ export class InstallTrackerSingleton
         return this.markInstallAsInUseWithInstallLock(installExePath, false, InstallTrackerSingleton.sessionId);
     }
 
+    /**
+     * Removes sessions from the persisted state whose IPC mutex is no longer held,
+     * indicating the owning VS Code process has exited or crashed.
+     * Should be called once on extension startup to avoid unbounded growth of the session map.
+     */
+    public async pruneStaleSessions(): Promise<void>
+    {
+        await executeWithLock(this.eventStream, false, this.getLockFilePathForKeySimple('installed'), 5, 200000,
+            async () =>
+            {
+                const serializedData = this.extensionState.get<Record<string, string[]>>(this.sessionInstallsKey, {});
+                const existingSessions = deserializeMapOfSets<string, string>(serializedData);
+                const initialCount = existingSessions.size;
+                let pruned = 0;
+
+                for (const sessionId of Array.from(existingSessions.keys()))
+                {
+                    if (sessionId === InstallTrackerSingleton.sessionId)
+                    {
+                        continue;
+                    }
+
+                    if (await this.isSessionDead(sessionId))
+                    {
+                        existingSessions.delete(sessionId);
+                        pruned++;
+                    }
+                }
+
+                if (pruned > 0)
+                {
+                    await this.extensionState.update(this.sessionInstallsKey, serializeMapOfSets(existingSessions));
+                    this.eventStream.post(new PrunedStaleSessions(`Pruned ${pruned} stale sessions out of ${initialCount}. Remaining: ${existingSessions.size}.`));
+                }
+            });
+    }
+
+    /**
+     * Removes the current session from the persisted state.
+     * Intended for clean shutdown via the extension's deactivate() export.
+     */
+    public async removeCurrentSession(): Promise<void>
+    {
+        await executeWithLock(this.eventStream, false, this.getLockFilePathForKeySimple('installed'), 5, 200000,
+            async () =>
+            {
+                const serializedData = this.extensionState.get<Record<string, string[]>>(this.sessionInstallsKey, {});
+                const existingSessions = deserializeMapOfSets<string, string>(serializedData);
+
+                if (existingSessions.delete(InstallTrackerSingleton.sessionId))
+                {
+                    await this.extensionState.update(this.sessionInstallsKey, serializeMapOfSets(existingSessions));
+                    this.eventStream.post(new RemovingCurrentSession(`Removed session ${InstallTrackerSingleton.sessionId} from state during deactivation.`));
+                }
+            });
+    }
+
+    /**
+     * Checks whether a session's IPC mutex can be acquired, meaning the process that held it has exited.
+     * @returns true if the session is dead (mutex was acquirable), false if still alive.
+     */
+    private async isSessionDead(sessionId: string): Promise<boolean>
+    {
+        const logger = new EventStreamNodeIPCMutexLoggerWrapper(this.eventStream, sessionId);
+        const mutex = new NodeIPCMutex(sessionId, logger, '');
+
+        // eslint-disable-next-line @typescript-eslint/require-await
+        return mutex.acquire(async () =>
+        {
+            this.eventStream.post(new DependentIsDead(`Session ${sessionId} is no longer live.`));
+            return true;
+        }, 10, 30, `${sessionId}-${crypto.randomUUID()}`).catch(() => false);
+    }
+
     protected markInstallAsInUseWithInstallLock(installExePath: string, alreadyHoldingLock: boolean, sessionId: string)
     {
         return executeWithLock(this.eventStream, alreadyHoldingLock, this.getLockFilePathForKeySimple('installed'), 5, 200000,

vscode-dotnet-runtime-library/src/EventStream/EventStreamEvents.ts

@@ -1242,6 +1242,16 @@ export class SearchingLiveDependents extends DotnetCustomMessageEvent
     public readonly eventName = 'SearchingLiveDependents';
 }
 
+export class PrunedStaleSessions extends DotnetCustomMessageEvent
+{
+    public readonly eventName = 'PrunedStaleSessions';
+}
+
+export class RemovingCurrentSession extends DotnetCustomMessageEvent
+{
+    public readonly eventName = 'RemovingCurrentSession';
+}
+
 export class CacheAliasCreated extends DotnetCustomMessageEvent
 {
     public readonly eventName = 'CacheAliasCreated';

vscode-dotnet-runtime-library/src/test/unit/InstallTracker.test.ts

@@ -775,4 +775,117 @@ suite('InstallTracker Session Mutex Tests', function ()
         }
     }).timeout(testTimeoutTime);
 
+    test('pruneStaleSessions removes dead sessions from state', async () =>
+    {
+        const extensionState = new MockExtensionContext();
+        const tracker = new MockInstallTracker(new MockEventStream(), extensionState);
+        const staleSession1 = generateRandomSessionId();
+        const staleSession2 = generateRandomSessionId();
+        const installExePath = path.join(getRandomValidFakeDir(), getDotnetExecutable());
+
+        try
+        {
+            // Seed two stale sessions (no process holding their mutex)
+            await tracker.markInstallAsInUseBySession(staleSession1, installExePath);
+            await tracker.markInstallAsInUseBySession(staleSession2, installExePath);
+
+            const storedBefore = extensionState.get<Record<string, string[]>>('dotnet.returnedInstallDirectories', {});
+            assert.isDefined(storedBefore[staleSession1], 'Stale session 1 should exist before prune');
+            assert.isDefined(storedBefore[staleSession2], 'Stale session 2 should exist before prune');
+
+            await tracker.pruneStaleSessions();
+
+            const storedAfter = extensionState.get<Record<string, string[]>>('dotnet.returnedInstallDirectories', {});
+            assert.isUndefined(storedAfter[staleSession1], 'Stale session 1 should be removed after prune');
+            assert.isUndefined(storedAfter[staleSession2], 'Stale session 2 should be removed after prune');
+        }
+        finally
+        {
+            await tracker.endAnySingletonTrackingSessions();
+        }
+    }).timeout(testTimeoutTime);
+
+    test('pruneStaleSessions preserves live sessions', async () =>
+    {
+        const extensionState = new MockExtensionContext();
+        const tracker = new MockInstallTracker(new MockEventStream(), extensionState);
+        const liveSessionId = generateRandomSessionId();
+        const staleSessionId = generateRandomSessionId();
+        const installExePath = path.join(getRandomValidFakeDir(), getDotnetExecutable());
+        const processes: { child: ChildProcess, sessionId: string }[] = [];
+
+        try
+        {
+            // Spawn a process that holds the mutex for the live session
+            const { child } = await spawnMutexHolderProcess(liveSessionId);
+            processes.push({ child, sessionId: liveSessionId });
+
+            await tracker.markInstallAsInUseBySession(liveSessionId, installExePath);
+            await tracker.markInstallAsInUseBySession(staleSessionId, installExePath);
+
+            await tracker.pruneStaleSessions();
+
+            const storedAfter = extensionState.get<Record<string, string[]>>('dotnet.returnedInstallDirectories', {});
+            assert.isDefined(storedAfter[liveSessionId], 'Live session should be preserved after prune');
+            assert.isUndefined(storedAfter[staleSessionId], 'Stale session should be removed after prune');
+        }
+        finally
+        {
+            await cleanupMutexHolders(processes);
+            await tracker.endAnySingletonTrackingSessions();
+        }
+    }).timeout(testTimeoutTime);
+
+    test('pruneStaleSessions does not remove the current session', async () =>
+    {
+        const extensionState = new MockExtensionContext();
+        const tracker = new MockInstallTracker(new MockEventStream(), extensionState);
+        const currentSessionId = tracker.getSessionId();
+        const installExePath = path.join(getRandomValidFakeDir(), getDotnetExecutable());
+
+        try
+        {
+            await tracker.markInstallAsInUseBySession(currentSessionId, installExePath);
+
+            await tracker.pruneStaleSessions();
+
+            const storedAfter = extensionState.get<Record<string, string[]>>('dotnet.returnedInstallDirectories', {});
+            assert.isDefined(storedAfter[currentSessionId], 'Current session should be preserved after prune');
+        }
+        finally
+        {
+            await tracker.endAnySingletonTrackingSessions();
+        }
+    }).timeout(testTimeoutTime);
+
+    test('removeCurrentSession removes only the current session from state', async () =>
+    {
+        const extensionState = new MockExtensionContext();
+        const tracker = new MockInstallTracker(new MockEventStream(), extensionState);
+        const currentSessionId = tracker.getSessionId();
+        const otherSessionId = generateRandomSessionId();
+        const installExePath = path.join(getRandomValidFakeDir(), getDotnetExecutable());
+        const processes: { child: ChildProcess, sessionId: string }[] = [];
+
+        try
+        {
+            const { child } = await spawnMutexHolderProcess(otherSessionId);
+            processes.push({ child, sessionId: otherSessionId });
+
+            await tracker.markInstallAsInUseBySession(currentSessionId, installExePath);
+            await tracker.markInstallAsInUseBySession(otherSessionId, installExePath);
+
+            await tracker.removeCurrentSession();
+
+            const storedAfter = extensionState.get<Record<string, string[]>>('dotnet.returnedInstallDirectories', {});
+            assert.isUndefined(storedAfter[currentSessionId], 'Current session should be removed after deactivation');
+            assert.isDefined(storedAfter[otherSessionId], 'Other session should be preserved');
+        }
+        finally
+        {
+            await cleanupMutexHolders(processes);
+            await tracker.endAnySingletonTrackingSessions();
+        }
+    }).timeout(testTimeoutTime);
+
 });