From f7f5ee780a69191d10b996966b5e2263d0bd0a95 Mon Sep 17 00:00:00 2001
From: Cory Rivera <corivera@microsoft.com>
Date: Fri, 22 Sep 2023 08:50:05 -0700
Subject: [PATCH] Use sync callback to stop Jupyter server when ADS process
 exits (#24501)

---
 extensions/notebook/src/common/utils.ts       |  4 +++
 .../notebook/src/jupyter/serverInstance.ts    | 28 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/extensions/notebook/src/common/utils.ts b/extensions/notebook/src/common/utils.ts
index 6cf9fd5b0c..10573a2103 100644
--- a/extensions/notebook/src/common/utils.ts
+++ b/extensions/notebook/src/common/utils.ts
@@ -52,6 +52,10 @@ export function executeBufferedCommand(cmd: string, options: childProcess.ExecOp
 	});
 }
 
+export function executeBufferedCommandSync(cmd: string, options: childProcess.ExecOptions): string {
+	return childProcess.execSync(cmd, options).toString();
+}
+
 export function executeStreamedCommand(cmd: string, options: childProcess.SpawnOptions, outputChannel?: vscode.OutputChannel): Thenable<void> {
 	return new Promise<void>((resolve, reject) => {
 		// Start the command
diff --git a/extensions/notebook/src/jupyter/serverInstance.ts b/extensions/notebook/src/jupyter/serverInstance.ts
index ebb1a61171..38fb49f23e 100644
--- a/extensions/notebook/src/jupyter/serverInstance.ts
+++ b/extensions/notebook/src/jupyter/serverInstance.ts
@@ -142,6 +142,30 @@ export class PerFolderServerInstance implements IServerInstance {
 		}
 	}
 
+	private stopSync(): void {
+		try {
+			this._isStopping = true;
+			if (this.baseDir) {
+				let exists = fs.pathExistsSync(this.baseDir);
+				if (exists) {
+					fs.removeSync(this.baseDir);
+				}
+			}
+			if (this._isStarted) {
+				let install = this.options.install;
+				let stopCommand = `"${install.pythonExecutable}" "${this.notebookScriptPath}" stop ${this._port}`;
+				utils.executeBufferedCommandSync(stopCommand, install.execOptions);
+			}
+		} 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;
+			ensureProcessEnded(this.childProcess);
+			this.handleConnectionClosed();
+		}
+	}
+
 	private async configureJupyter(): Promise<void> {
 		this.createInstanceFolders();
 		let resourcesFolder = path.join(this.options.install.extensionPath, 'resources', constants.jupyterConfigRootFolder);
@@ -269,7 +293,9 @@ export class PerFolderServerInstance implements IServerInstance {
 		this.childProcess.addListener('error', this.handleConnectionError);
 		this.childProcess.addListener('exit', this.handleConnectionClosed);
 
-		process.addListener('exit', this.stop);
+		// Have to use the synchronous Stop callback for Exit events since it's required in the Node.JS docs: https://nodejs.org/api/process.html#event-exit
+		// The process exits immediately after calling the Stop callback, so if execution is yielded with an await, then any code after that isn't run.
+		process.addListener('exit', this.stopSync);
 
 		// 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
 	}