Improve Cleanup of Jupyter processes on Notebook and/or ADS Close (#5142)

* Close jupyter and python

* Ensure we stop jupyter correctly on process end

* dont stopServer from clientSession shutdown

* PR comments

* close notebook after each test

extensions/integration-tests/src/notebook.test.ts

@@ -20,7 +20,7 @@ if (context.RunTest) {
 		setup(function () {
 			console.log(`Start "${this.currentTest.title}"`);
 		});
-		teardown(function () {
+		teardown(async function () {
 			let testName = this.currentTest.title;
 			try {
 				let fileName = getFileName(testName);
@@ -28,6 +28,7 @@ if (context.RunTest) {
 					fs.unlinkSync(fileName);
 					console.log(`"${fileName}" is deleted.`);
 				}
+				await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
 			}
 			catch (err) {
 				console.log(err);

extensions/notebook/src/jupyter/serverInstance.ts

@@ -76,19 +76,20 @@ export class ServerInstanceUtils {
 		return spawn(command, args, options);
 	}
 
-	public checkProcessDied(childProcess: ChildProcess): void {
+	public ensureProcessEnded(childProcess: ChildProcess): void {
 		if (!childProcess) {
 			return;
 		}
-		// Wait 10 seconds and then force kill. Jupyter stop is slow so this seems a reasonable time limit
+		// Wait 5 seconds and then force kill. Jupyter stop is slow so this seems a reasonable time limit
 		setTimeout(() => {
 			// Test if the process is still alive. Throws an exception if not
 			try {
-				process.kill(childProcess.pid, <any>0);
+				process.kill(childProcess.pid, 'SIGKILL');
 			} catch (error) {
+				console.log(error);
 				// All is fine.
 			}
-		}, 10000);
+		}, 5000);
 	}
 }
 
@@ -156,13 +157,14 @@ export class PerNotebookServerInstance implements IServerInstance {
 				let install = this.options.install;
 				let stopCommand = `"${install.pythonExecutable}" -m jupyter notebook stop ${this._port}`;
 				await this.utils.executeBufferedCommand(stopCommand, install.execOptions, install.outputChannel);
-				this._isStarted = false;
-				this.utils.checkProcessDied(this.childProcess);
-				this.handleConnectionClosed();
 			}
 		} catch (error) {
 			// For now, we don't care as this is non-critical
 			this.notify(this.options.install, localize('serverStopError', 'Error stopping Notebook Server: {0}', utils.getErrorMessage(error)));
+		} finally {
+			this._isStarted = false;
+			this.utils.ensureProcessEnded(this.childProcess);
+			this.handleConnectionClosed();
 		}
 	}
 
@@ -290,6 +292,8 @@ export class PerNotebookServerInstance implements IServerInstance {
 		this.childProcess.addListener('error', this.handleConnectionError);
 		this.childProcess.addListener('exit', this.handleConnectionClosed);
 
+		process.addListener('exit', this.stop);
+
 		// TODO #897 covers serializing stdout and stderr to a location where we can read from so that user can see if they run into trouble
 	}
 
@@ -363,9 +367,10 @@ export class PerNotebookServerInstance implements IServerInstance {
 		env['MSHOST_ENVIRONMENT'] = 'ADSClient-' + vscode.version;
 
 		// Start the notebook process
-		let options = {
+		let options: SpawnOptions = {
 			shell: true,
-			env: env
+			env: env,
+			detached: false
 		};
 		let childProcess = this.utils.spawn(startCommand, [], options);
 		return childProcess;

extensions/notebook/src/test/model/serverInstance.test.ts

@@ -38,7 +38,7 @@ describe('Jupyter server instance', function (): void {
 		mockInstall.setup(i => i.outputChannel).returns(() => mockOutputChannel.object);
 		mockInstall.setup(i => i.pythonExecutable).returns(() => 'python3');
 		mockUtils = TypeMoq.Mock.ofType(ServerInstanceUtils);
-		mockUtils.setup(u => u.checkProcessDied(TypeMoq.It.isAny())).returns(() => undefined);
+		mockUtils.setup(u => u.ensureProcessEnded(TypeMoq.It.isAny())).returns(() => undefined);
 		serverInstance = new PerNotebookServerInstance({
 			documentPath: expectedPath,
 			install: mockInstall.object

src/sql/workbench/api/node/extHostNotebook.ts

@@ -50,7 +50,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		let manager = this.findManagerForUri(uriString);
 		if (manager) {
 			manager.provider.handleNotebookClosed(uri);
-			// Note: deliberately not removing handle.
+			this._adapters.delete(manager.handle);
 		}
 	}
 
@@ -249,7 +249,6 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 	private _createDisposable(handle: number): Disposable {
 		return new Disposable(() => {
 			this._adapters.delete(handle);
-			this._proxy.$unregisterNotebookProvider(handle);
 		});
 	}
 

src/sql/workbench/parts/notebook/models/notebookModel.ts

@@ -777,7 +777,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			}
 			await this.shutdownActiveSession();
 		} catch (err) {
-			this.notifyError(localize('shutdownError', "An error occurred when closing the notebook: {0}", notebookUtils.getErrorMessage(err)));
+			console.log('An error occurred when closing the notebook: {0}', notebookUtils.getErrorMessage(err));
 		}
 	}
 

src/sqltest/workbench/api/exthostNotebook.test.ts

@@ -111,10 +111,10 @@ suite('ExtHostNotebook Tests', () => {
 
 		});
 
-		test('Should call unregister on disposing', () => {
+		test('Should not call unregister on disposing', () => {
 			let disposable = extHostNotebook.registerNotebookProvider(notebookProviderMock.object);
 			disposable.dispose();
-			mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.once());
+			mockProxy.verify(p => p.$unregisterNotebookProvider(TypeMoq.It.isValue(savedHandle)), TypeMoq.Times.never());
 		});
 	});
 });