Store state correctly and Remove old state

microsoft · May 24, 2023 · 4455605 · 4455605
1 parent 1bbf9ba
commit 4455605
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 11 deletions.
diff --git a/src/kernels/execution/lastCellExecutionTracker.ts b/src/kernels/execution/lastCellExecutionTracker.ts
@@ -4,7 +4,7 @@
 import { injectable, inject, named } from 'inversify';
 import { IDisposable, IDisposableRegistry, IMemento, WORKSPACE_MEMENTO } from '../../platform/common/types';
 import { Disposables } from '../../platform/common/utils';
-import { IKernel, ResumeCellExecutionInformation, isRemoteConnection } from '../types';
+import { IKernel, RemoteKernelConnectionMetadata, ResumeCellExecutionInformation, isRemoteConnection } from '../types';
 import type { KernelMessage } from '@jupyterlab/services';
 import { IAnyMessageArgs } from '@jupyterlab/services/lib/kernel/kernel';
 import { disposeAllDisposables } from '../../platform/common/helpers';
@@ -13,8 +13,12 @@ import { noop, swallowExceptions } from '../../platform/common/utils/misc';
 import { getParentHeaderMsgId } from './cellExecutionMessageHandler';
 import { IJupyterServerUriEntry, IJupyterServerUriStorage, JupyterServerProviderHandle } from '../jupyter/types';
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
-import { computeServerId } from '../jupyter/jupyterUtils';
+import { jupyterServerHandleToString } from '../jupyter/jupyterUtils';
 
+const LAST_EXECUTED_CELL_PREFIX = `LAST_EXECUTED_CELL_V2_`;
+function getConnectionStatePrefix(provider: JupyterServerProviderHandle) {
+    return `${LAST_EXECUTED_CELL_PREFIX}${jupyterServerHandleToString(provider)}`;
+}
 type CellExecutionInfo = Omit<ResumeCellExecutionInformation, 'token'> & { kernelId: string; cellIndex: number };
 /**
  * Keeps track of the last cell that was executed for a notebook along with the time and execution count.
@@ -34,8 +38,8 @@ export class LastCellExecutionTracker extends Disposables implements IExtensionS
     public activate(): void {
         this.serverStorage.onDidRemove(this.onDidRemoveServerUris, this, this.disposables);
     }
-    private async getStateKey(serverHandle: JupyterServerProviderHandle) {
-        return `LAST_EXECUTED_CELL_${await computeServerId(serverHandle)}`;
+    private async getStateKey(connection: RemoteKernelConnectionMetadata) {
+        return `${getConnectionStatePrefix(connection.serverHandle)}${connection.id}`;
     }
     public async getLastTrackedCellExecution(
         notebook: NotebookDocument,
@@ -48,7 +52,7 @@ export class LastCellExecutionTracker extends Disposables implements IExtensionS
             return;
         }
         const data = this.workspaceMemento.get<{ [key: string]: CellExecutionInfo }>(
-            await this.getStateKey(kernel.kernelConnectionMetadata.serverHandle),
+            await this.getStateKey(kernel.kernelConnectionMetadata),
             {}
         );
         return data[notebook.uri.toString()];
@@ -134,7 +138,7 @@ export class LastCellExecutionTracker extends Disposables implements IExtensionS
             return;
         }
 
-        const id = await this.getStateKey(kernel.kernelConnectionMetadata.serverHandle);
+        const id = await this.getStateKey(kernel.kernelConnectionMetadata);
         this.chainedPromises = this.chainedPromises.finally(() => {
             const notebookId = cell.notebook.uri.toString();
             const currentState = this.workspaceMemento.get<{ [key: string]: Partial<CellExecutionInfo> }>(id, {});
@@ -152,7 +156,7 @@ export class LastCellExecutionTracker extends Disposables implements IExtensionS
             return;
         }
 
-        const id = await this.getStateKey(kernel.kernelConnectionMetadata.serverHandle);
+        const id = await this.getStateKey(kernel.kernelConnectionMetadata);
         this.chainedPromises = this.chainedPromises.finally(() => {
             const notebookId = cell.notebook.uri.toString();
             const currentState = this.workspaceMemento.get<{ [key: string]: Partial<CellExecutionInfo> }>(id, {});
@@ -163,9 +167,11 @@ export class LastCellExecutionTracker extends Disposables implements IExtensionS
     private onDidRemoveServerUris(removedServers: IJupyterServerUriEntry[]) {
         this.chainedPromises = this.chainedPromises.finally(() =>
             Promise.all(
-                removedServers
-                    .map((item) => this.getStateKey(item.serverHandle))
-                    .map(async (id) => this.workspaceMemento.update(await id, undefined).then(noop, noop))
+                removedServers.map(async (id) => {
+                    const prefixOfKeysToRemove = getConnectionStatePrefix(id.serverHandle);
+                    const keys = this.workspaceMemento.keys().filter((k) => k.startsWith(prefixOfKeysToRemove));
+                    await Promise.all(keys.map((key) => this.workspaceMemento.update(key, undefined).then(noop, noop)));
+                })
             )
         );
     }

diff --git a/src/platform/common/cache.ts b/src/platform/common/cache.ts
@@ -5,7 +5,7 @@ import { Memento } from 'vscode';
 import { noop } from './utils/misc';
 import { IExtensionSyncActivationService } from '../activation/types';
 import { IApplicationEnvironment, IWorkspaceService } from './application/types';
-import { GLOBAL_MEMENTO, ICryptoUtils, IMemento } from './types';
+import { GLOBAL_MEMENTO, ICryptoUtils, IMemento, WORKSPACE_MEMENTO } from './types';
 import { inject, injectable, named } from 'inversify';
 import { getFilePath } from './platform/fs-paths';
 
@@ -14,11 +14,13 @@ export class OldCacheCleaner implements IExtensionSyncActivationService {
     constructor(
         @inject(IWorkspaceService) private readonly workspace: IWorkspaceService,
         @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly globalState: Memento,
+        @inject(IMemento) @named(WORKSPACE_MEMENTO) private readonly workspaceState: Memento,
         @inject(ICryptoUtils) private readonly crypto: ICryptoUtils,
         @inject(IApplicationEnvironment) private readonly appEnv: IApplicationEnvironment
     ) {}
     public activate(): void {
         this.removeOldCachedItems().then(noop, noop);
+        this.removeOldWorkspaceCachedItems().then(noop, noop);
     }
     async removeOldCachedItems(): Promise<void> {
         await Promise.all(
@@ -42,6 +44,12 @@ export class OldCacheCleaner implements IExtensionSyncActivationService {
                 .map((key) => this.globalState.update(key, undefined).then(noop, noop))
         );
     }
+    async removeOldWorkspaceCachedItems(): Promise<void> {
+        const keys = this.workspaceState
+            .keys()
+            .filter((k) => k.startsWith('LAST_EXECUTED_CELL_') && !k.startsWith('LAST_EXECUTED_CELL_V2_'));
+        await Promise.all(keys.map((key) => this.workspaceState.update(key, undefined).then(noop, noop)));
+    }
 
     async getUriAccountKey(): Promise<string> {
         if (this.workspace.rootFolder) {