Starting Fewer Jupyter Servers for Notebooks (#7744)

* Start fewer Jupyter servers

* Windows fix for drive casing

* PR Feedback

* Quick fix

* Fixing bug

* Ensure environment variables set 4 session startup

* test fix

* Dummy commit to update comment

extensions/notebook/src/jupyter/jupyterNotebookProvider.ts

@@ -3,16 +3,17 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-'use strict';
-
 import { nb } from 'azdata';
 import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
+import * as path from 'path';
 const localize = nls.loadMessageBundle();
 
 import * as constants from '../common/constants';
+import * as utils from '../common/utils';
 import { JupyterNotebookManager } from './jupyterNotebookManager';
 import { LocalJupyterServerManager } from './jupyterServerManager';
+import { JupyterSessionManager } from './jupyterSessionManager';
 
 export type ServerManagerFactory = (documentUri: vscode.Uri) => LocalJupyterServerManager;
 
@@ -31,12 +32,16 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
 	}
 
 	private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager {
-		let uriString = notebookUri.toString();
-		let manager = this.managerTracker.get(uriString);
+		let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
+		let manager = this.managerTracker.get(baseFolder);
 		if (!manager) {
-			let serverManager = this.createServerManager(notebookUri);
+			let baseFolderUri = vscode.Uri.file(baseFolder);
+			if (!baseFolderUri) {
+				baseFolderUri = notebookUri;
+			}
+			let serverManager = this.createServerManager(baseFolderUri);
 			manager = new JupyterNotebookManager(serverManager);
-			this.managerTracker.set(uriString, manager);
+			this.managerTracker.set(baseFolder, manager);
 		}
 		return manager;
 	}
@@ -46,15 +51,52 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
 			// As this is a notification method, will skip throwing an error here
 			return;
 		}
-		let uriString = notebookUri.toString();
-		let manager = this.managerTracker.get(uriString);
+		let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
+		let manager = this.managerTracker.get(baseFolder);
 		if (manager) {
-			this.managerTracker.delete(uriString);
-			manager.dispose();
+			let sessionManager = (manager.sessionManager as JupyterSessionManager);
+			let session = sessionManager.listRunning().find(e => e.path === notebookUri.fsPath);
+			if (session) {
+				manager.sessionManager.shutdown(session.id);
+			}
+			if (sessionManager.listRunning().length === 0) {
+				const FiveMinutesInMs = 5 * 60 * 1000;
+				setTimeout(() => {
+					if (sessionManager.listRunning().length === 0) {
+						this.managerTracker.delete(baseFolder);
+						manager.dispose();
+					}
+				}, FiveMinutesInMs);
+			}
 		}
 	}
 
 	public get standardKernels(): nb.IStandardKernel[] {
 		return [];
 	}
+
+	private transformToBaseFolder(notebookPath: string): string {
+		let parsedPath = path.parse(notebookPath);
+		let userHome = utils.getUserHome();
+		let relativePathStrUserHome = path.relative(notebookPath, userHome);
+		if (notebookPath === '.' || relativePathStrUserHome.endsWith('..') || relativePathStrUserHome === '') {
+			// If you don't match the notebookPath's casing for drive letters,
+			// a 404 will result when trying to create a new session on Windows
+			if (notebookPath && userHome && notebookPath[0].toLowerCase() === userHome[0].toLowerCase()) {
+				userHome = notebookPath[0] + userHome.substr(1);
+			}
+			// If the user is using a system version of python, then
+			// '.' will try to create a notebook in a system directory.
+			// Since this will fail due to permissions issues, use the user's
+			// home folder instead.
+			return userHome;
+		} else {
+			let splitDirName: string[] = path.dirname(notebookPath).split(path.sep);
+			if (splitDirName.length > 1) {
+				return path.join(parsedPath.root, splitDirName[1]);
+			} else {
+				return userHome;
+			}
+		}
+	}
 }