Wait for up to 6s for extension hosts to exit, then proceed to kill them

src/vs/code/electron-main/app.ts

@@ -45,7 +45,7 @@ import { resolveShellEnv } from 'vs/platform/environment/node/shellEnv';
 import { IExtensionUrlTrustService } from 'vs/platform/extensionManagement/common/extensionUrlTrust';
 import { ExtensionUrlTrustService } from 'vs/platform/extensionManagement/node/extensionUrlTrustService';
 import { IExtensionHostStarter, ipcExtensionHostStarterChannelName } from 'vs/platform/extensions/common/extensionHostStarter';
-import { ExtensionHostStarter } from 'vs/platform/extensions/node/extensionHostStarter';
+import { DirectMainProcessExtensionHostStarter } from 'vs/platform/extensions/electron-main/directMainProcessExtensionHostStarter';
 import { IExternalTerminalMainService } from 'vs/platform/externalTerminal/common/externalTerminal';
 import { LinuxExternalTerminalService, MacExternalTerminalService, WindowsExternalTerminalService } from 'vs/platform/externalTerminal/node/externalTerminalService';
 import { IFileService } from 'vs/platform/files/common/files';
@@ -517,7 +517,7 @@ export class CodeApplication extends Disposable {
 		services.set(IExtensionUrlTrustService, new SyncDescriptor(ExtensionUrlTrustService));
 
 		// Extension Host Starter
-		services.set(IExtensionHostStarter, new SyncDescriptor(ExtensionHostStarter));
+		services.set(IExtensionHostStarter, new SyncDescriptor(DirectMainProcessExtensionHostStarter));
 
 		// Storage
 		services.set(IStorageMainService, new SyncDescriptor(StorageMainService));

src/vs/platform/extensions/electron-main/directMainProcessExtensionHostStarter.ts

@@ -0,0 +1,27 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ExtensionHostStarter, IPartialLogService } from 'vs/platform/extensions/node/extensionHostStarter';
+import { ILifecycleMainService } from 'vs/platform/lifecycle/electron-main/lifecycleMainService';
+import { ILogService } from 'vs/platform/log/common/log';
+
+export class DirectMainProcessExtensionHostStarter extends ExtensionHostStarter {
+
+	constructor(
+		@ILogService logService: IPartialLogService,
+		@ILifecycleMainService lifecycleMainService: ILifecycleMainService
+	) {
+		super(logService);
+
+		lifecycleMainService.onWillShutdown((e) => {
+			const exitPromises: Promise<void>[] = [];
+			for (const [, extHost] of this._extHosts) {
+				exitPromises.push(extHost.waitForExit(6000));
+			}
+			e.join(Promise.all(exitPromises).then(() => { }));
+		});
+	}
+
+}

src/vs/platform/extensions/node/extensionHostStarter.ts

@@ -16,6 +16,7 @@ import { registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { mixin } from 'vs/base/common/objects';
 import { cwd } from 'vs/base/common/process';
 import { StopWatch } from 'vs/base/common/stopwatch';
+import { timeout } from 'vs/base/common/async';
 
 export interface IPartialLogService {
 	readonly _serviceBrand: undefined;
@@ -40,6 +41,7 @@ class ExtensionHostProcess extends Disposable {
 	readonly onExit = this._onExit.event;
 
 	private _process: ChildProcess | null = null;
+	private _hasExited: boolean = false;
 
 	constructor(
 		public readonly id: string,
@@ -48,10 +50,6 @@ class ExtensionHostProcess extends Disposable {
 		super();
 	}
 
-	register(disposable: IDisposable) {
-		this._register(disposable);
-	}
-
 	start(opts: IExtensionHostProcessOptions): { pid: number; } {
 		const sw = StopWatch.create(false);
 		opts.env = opts.env || {};
@@ -92,6 +90,7 @@ class ExtensionHostProcess extends Disposable {
 		});
 
 		this._process.on('exit', (code: number, signal: string) => {
+			this._hasExited = true;
 			this._onExit.fire({ pid, code, signal });
 		});
 
@@ -130,14 +129,29 @@ class ExtensionHostProcess extends Disposable {
 		this._logService.info(`Killing extension host with pid ${this._process.pid}.`);
 		this._process.kill();
 	}
+
+	async waitForExit(maxWaitTimeMs: number): Promise<void> {
+		if (!this._process) {
+			return;
+		}
+		const pid = this._process.pid;
+		this._logService.info(`Waiting for extension host with pid ${pid} to exit.`);
+		await Promise.race([Event.toPromise(this.onExit), timeout(maxWaitTimeMs)]);
+
+		if (!this._hasExited) {
+			// looks like we timed out
+			this._logService.info(`Extension host with pid ${pid} did not exit within ${maxWaitTimeMs}ms.`);
+			this._process.kill();
+		}
+	}
 }
 
 export class ExtensionHostStarter implements IDisposable, IExtensionHostStarter {
 	_serviceBrand: undefined;
 
 	private static _lastId: number = 0;
 
-	private readonly _extHosts: Map<string, ExtensionHostProcess>;
+	protected readonly _extHosts: Map<string, ExtensionHostProcess>;
 
 	constructor(
 		@ILogService private readonly _logService: IPartialLogService