Kernel MRU for web and desktop (#11969)

* Kernel MRU for web and desktop

* oops

src/notebooks/controllers/connectionMru.node.ts

@@ -11,17 +11,18 @@ import { getNotebookMetadata } from '../../platform/common/utils';
 import { swallowExceptions } from '../../platform/common/utils/decorators';
 import { traceWarning } from '../../platform/logging';
 import { IKernelRankingHelper, IConnectionMru } from './types';
-import { EOL } from 'os';
 import { getDisplayPath } from '../../platform/common/platform/fs-paths';
 import { createDeferredFromPromise } from '../../platform/common/utils/async';
 import { IWorkspaceService } from '../../platform/common/application/types';
 import { getTelemetrySafeHashedString } from '../../platform/telemetry/helpers';
+import * as fs from 'fs';
+import { MaxMRUSizePerNotebook, MRUItem } from './connectionMru';
 
 const MruFolder = 'notebook-connection-mru';
 @injectable()
-export class ConnectionMRU implements IConnectionMru {
+export class ConnectionMru implements IConnectionMru {
     private documentSourceMapping = new WeakMap<NotebookDocument, Set<KernelConnectionMetadata>>();
-    private documentMruContents = new WeakMap<NotebookDocument, Promise<string | undefined>>();
+    private documentMruContents = new WeakMap<NotebookDocument, Promise<MRUItem[]>>();
     private notebookCacheFileName = new WeakMap<NotebookDocument, Promise<Uri>>();
 
     constructor(
@@ -48,30 +49,47 @@ export class ConnectionMRU implements IConnectionMru {
         const file = await this.getCacheFileName(notebook);
         if (!(await this.fs.exists(path.dirname(file)))) {
             try {
-                await this.fs.createDirectory(this.context.storageUri);
+                await this.fs.createDirectory(path.dirname(file));
             } catch (ex) {
                 traceWarning(`Failed to create directory ${getDisplayPath(this.context.storageUri)}`, ex);
             }
         }
         if (await this.fs.exists(path.dirname(file))) {
-            const [connectionIdsInMru, connectionIds] = await Promise.all([
-                this.fs
-                    .readFile(file)
-                    .then((c) => c.splitLines({ trim: true, removeEmptyEntries: true }))
-                    .catch(() => <string[]>[]),
-                Promise.all(Array.from(connections).map((item) => item.getHashId()))
+            const connectionIdsUsedInSession = new Set<string>();
+            const [currentMru, connectionsUsedInSession] = await Promise.all([
+                this.loadNotebookCache(notebook),
+                Promise.all(
+                    Array.from(connections).map(async (item) => {
+                        const id = await item.getHashId();
+                        connectionIdsUsedInSession.add(id);
+                        return <MRUItem>[Date.now(), id];
+                    })
+                )
             ]);
-            const updatedConnectionIds = Array.from(new Set(connectionIdsInMru.concat(connectionIds)));
-            const newContents = updatedConnectionIds.join(EOL);
-            this.documentMruContents.set(notebook, Promise.resolve(newContents));
-            await this.fs.writeFile(file, newContents);
+            const idsInMru = new Set<string>();
+
+            // For existing items, update the last used time.
+            currentMru.forEach((mru) => {
+                idsInMru.add(mru[1]);
+                if (connectionIdsUsedInSession.has(mru[1])) {
+                    mru[0] = Date.now();
+                }
+            });
+            // If we have more than 10 items, then remove the oldest items.
+            const newMruItems = currentMru
+                .concat(connectionsUsedInSession.filter((item) => !idsInMru.has(item[1])))
+                .sort((a, b) => b[0] - a[0])
+                .slice(0, MaxMRUSizePerNotebook);
+
+            this.documentMruContents.set(notebook, Promise.resolve(newMruItems));
+            await this.fs.writeFile(file, JSON.stringify(newMruItems));
         }
     }
     public async clear(): Promise<void> {
         const cacheFolder = Uri.joinPath(this.context.globalStorageUri, MruFolder);
-        await this.fs
-            .delete(cacheFolder)
-            .catch((ex) => traceWarning(`Failed to delete MRU cache folder ${getDisplayPath(cacheFolder)}`, ex));
+        new Promise<void>((resolve, reject) =>
+            fs.rmdir(cacheFolder.fsPath, { recursive: true }, (ex) => (ex ? reject(ex) : resolve()))
+        ).catch((ex) => traceWarning(`Failed to delete MRU cache folder ${getDisplayPath(cacheFolder)}`, ex));
     }
     /**
      * Checks whether a connection was used by a notebook.
@@ -103,10 +121,10 @@ export class ConnectionMRU implements IConnectionMru {
                 const workspaceId = this.workspace.getWorkspaceFolderIdentifier(notebook.uri, 'global');
                 const workspaceHash = await getTelemetrySafeHashedString(workspaceId);
                 return Uri.joinPath(
-                    this.context.globalStorageUri,
+                    this.context.storageUri || this.context.globalStorageUri,
                     MruFolder,
                     workspaceHash,
-                    `${path.basename(notebook.uri)}.last_used_connections.txt`
+                    `${path.basename(notebook.uri)}.last_used_connections.json`
                 );
             })();
             this.notebookCacheFileName.set(notebook, promise);
@@ -117,12 +135,18 @@ export class ConnectionMRU implements IConnectionMru {
         const mruFileContents = this.documentMruContents.get(notebook);
         if (!mruFileContents) {
             const promise = (async () => {
-                const file = await this.getCacheFileName(notebook);
-                return this.fs.readFile(file);
+                try {
+                    const file = await this.getCacheFileName(notebook);
+                    const contents = await this.fs.readFile(file);
+                    return JSON.parse(contents) as MRUItem[];
+                } catch {
+                    // File doesn't exist.
+                    return [];
+                }
             })();
             this.documentMruContents.set(notebook, promise);
         }
-        return this.documentMruContents.get(notebook);
+        return this.documentMruContents.get(notebook)!;
     }
     private async existsInFile(notebook: NotebookDocument, connection: KernelConnectionMetadata) {
         const connections = this.documentSourceMapping.get(notebook);
@@ -140,7 +164,7 @@ export class ConnectionMRU implements IConnectionMru {
                 this.loadNotebookCache(notebook),
                 connection.getHashId()
             ]);
-            return (contents || '').includes(connectionIdHash);
+            return contents.some((item) => item[1] === connectionIdHash);
         } catch {
             return false;
         }

src/notebooks/controllers/connectionMru.node.unit.test.ts

@@ -0,0 +1,166 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { assert } from 'chai';
+import { anything, capture, instance, mock, verify, when } from 'ts-mockito';
+import { Uri } from 'vscode';
+import { LocalKernelSpecConnectionMetadata, PythonKernelConnectionMetadata } from '../../kernels/types';
+import { IWorkspaceService } from '../../platform/common/application/types';
+import { IFileSystem } from '../../platform/common/platform/types';
+import { IExtensionContext } from '../../platform/common/types';
+import { TestNotebookDocument } from '../../test/datascience/notebook/executionHelper';
+import { ConnectionMru } from './connectionMru.node';
+import { IKernelRankingHelper } from './types';
+
+suite('Connection MRU (node)', () => {
+    let mru: ConnectionMru;
+    let rankingHelper: IKernelRankingHelper;
+    let context: IExtensionContext;
+    let fs: IFileSystem;
+    let workspace: IWorkspaceService;
+    const javaKernelSpec = LocalKernelSpecConnectionMetadata.create({
+        id: 'java',
+        kernelSpec: {
+            argv: ['java'],
+            display_name: 'java',
+            executable: 'java',
+            name: 'java',
+            language: 'java'
+        }
+    });
+    const pythonKernelSpec = PythonKernelConnectionMetadata.create({
+        id: 'python',
+        interpreter: {
+            id: 'python',
+            sysPrefix: '',
+            uri: Uri.file('python')
+        },
+        kernelSpec: {
+            argv: ['python'],
+            display_name: 'python',
+            executable: 'python',
+            name: 'python'
+        }
+    });
+    const notebook = new TestNotebookDocument(Uri.file('notebook1.ipynb'));
+    setup(() => {
+        rankingHelper = mock<IKernelRankingHelper>();
+        context = mock<IExtensionContext>();
+        fs = mock<IFileSystem>();
+        workspace = mock<IWorkspaceService>();
+        when(fs.exists(anything())).thenResolve(false);
+        when(fs.createDirectory(anything())).thenResolve();
+        when(context.storageUri).thenReturn(Uri.file('workspaceStorage'));
+        when(workspace.getWorkspaceFolderIdentifier(anything())).thenReturn('workspace1');
+        when(workspace.getWorkspaceFolderIdentifier(anything(), anything())).thenReturn('workspace1');
+        when(fs.readFile(anything())).thenReject(new Error('File not found'));
+        when(rankingHelper.isExactMatch(anything(), anything(), anything())).thenResolve(false);
+        mru = new ConnectionMru(instance(rankingHelper), instance(context), instance(fs), instance(workspace));
+    });
+
+    test('No MRU items for first time users', async () => {
+        const exists = await mru.exists(notebook, pythonKernelSpec);
+
+        assert.isFalse(exists);
+    });
+    test('Update MRU', async () => {
+        await mru.add(notebook, pythonKernelSpec);
+        const exists = await mru.exists(notebook, pythonKernelSpec);
+
+        assert.isTrue(exists);
+    });
+    test('Update file when updating MRU', async () => {
+        when(fs.exists(anything())).thenResolve(true);
+        assert.isFalse(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+
+        await mru.add(notebook, pythonKernelSpec);
+
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+        verify(fs.writeFile(anything(), anything())).atLeast(1);
+        let json = JSON.parse(capture(fs.writeFile).first()[1] as string) as [number, string][];
+        assert.strictEqual(json[0][1], await pythonKernelSpec.getHashId());
+
+        await mru.add(notebook, javaKernelSpec);
+
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook, javaKernelSpec));
+        json = JSON.parse(capture(fs.writeFile).second()[1] as string) as [number, string][];
+        assert.strictEqual(json[0][1], await pythonKernelSpec.getHashId());
+        assert.strictEqual(json[1][1], await javaKernelSpec.getHashId());
+    });
+    test('Load MRU from file', async () => {
+        when(fs.exists(anything())).thenResolve(true);
+        when(fs.readFile(anything())).thenResolve(JSON.stringify([[1, await pythonKernelSpec.getHashId()]]));
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+    });
+    test('Load MRU from file (with more than one item)', async () => {
+        when(fs.exists(anything())).thenResolve(true);
+        when(fs.readFile(anything())).thenResolve(
+            JSON.stringify([
+                [1, await pythonKernelSpec.getHashId()],
+                [2, await javaKernelSpec.getHashId()]
+            ])
+        );
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook, javaKernelSpec));
+    });
+    test('Update existing MRUs', async () => {
+        when(fs.exists(anything())).thenResolve(true);
+        when(fs.readFile(anything())).thenResolve(
+            JSON.stringify([
+                [1, await pythonKernelSpec.getHashId()],
+                [2, await javaKernelSpec.getHashId()]
+            ])
+        );
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook, javaKernelSpec));
+
+        await mru.add(notebook, pythonKernelSpec);
+
+        verify(fs.writeFile(anything(), anything())).atLeast(1);
+        let json = JSON.parse(capture(fs.writeFile).first()[1] as string) as [number, string][];
+        assert.strictEqual(json[0][1], await pythonKernelSpec.getHashId());
+        assert.isAtLeast(json[0][0], Date.now() - 60 * 1000);
+        assert.strictEqual(json[1][0], 2);
+
+        await mru.add(notebook, javaKernelSpec);
+
+        json = JSON.parse(capture(fs.writeFile).second()[1] as string) as [number, string][];
+        assert.strictEqual(json[0][1], await pythonKernelSpec.getHashId());
+        assert.strictEqual(json[1][1], await javaKernelSpec.getHashId());
+        assert.isAtLeast(json[1][0], Date.now() - 60 * 1000);
+    });
+    test('Use different MRUs per notebook document', async () => {
+        when(fs.exists(anything())).thenResolve(true);
+        const notebook2 = new TestNotebookDocument(Uri.file('notebook2.ipynb'));
+
+        assert.isFalse(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+        assert.isFalse(await mru.exists(notebook2, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook2, javaKernelSpec));
+
+        await mru.add(notebook, pythonKernelSpec);
+
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+        assert.isFalse(await mru.exists(notebook2, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook2, javaKernelSpec));
+
+        await mru.add(notebook2, javaKernelSpec);
+
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isFalse(await mru.exists(notebook, javaKernelSpec));
+        assert.isFalse(await mru.exists(notebook2, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook2, javaKernelSpec));
+
+        await Promise.all([mru.add(notebook, javaKernelSpec), mru.add(notebook2, pythonKernelSpec)]);
+
+        assert.isTrue(await mru.exists(notebook, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook, javaKernelSpec));
+        assert.isTrue(await mru.exists(notebook2, pythonKernelSpec));
+        assert.isTrue(await mru.exists(notebook2, javaKernelSpec));
+    });
+});

src/notebooks/controllers/connectionMru.ts

@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+export const MaxMRUSizePerNotebook = 10;
+// Keep the date when a connection was last used, we might need this, after all its an MRU
+export type MRUItem = [lastUsed: number, connectionId: string];