remote: gracefully disconnect from remote EH (#223422)

* remote: gracefully disconnect from remote EH

Addresses #211462 (comment)

This applies the same fix seen in #218972 to the extension host. It adds
a `disconnect` method that is called and awaited during the lifecycle
shutdown, which is implemented by th remote extension host connection.

The disconnection is awaited, but I think this is safe because it only
awaits `.drain()` and does so in the context of methods which were
already async at higher levels.

* fix tests

* Revert "cli: honor --connection-token-file in serve-web (#219041)"

This reverts commit d8ddc28. That
commit was incorrectly merged in the branch.

src/vs/workbench/services/extensions/browser/extensionService.ts

@@ -186,7 +186,7 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 
 	protected async _onExtensionHostExit(code: number): Promise<void> {
 		// Dispose everything associated with the extension host
-		this._doStopExtensionHosts();
+		await this._doStopExtensionHosts();
 
 		// If we are running extension tests, forward logs and exit code
 		const automatedWindow = mainWindow as unknown as IAutomatedWindow;

src/vs/workbench/services/extensions/common/abstractExtensionService.ts

@@ -197,23 +197,25 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 
 		this._register(this._lifecycleService.onWillShutdown(event => {
 			if (this._remoteAgentService.getConnection()) {
-				event.join(() => this._remoteAgentService.endConnection(), {
+				event.join(async () => {
+					// We need to disconnect the management connection before killing the local extension host.
+					// Otherwise, the local extension host might terminate the underlying tunnel before the
+					// management connection has a chance to send its disconnection message.
+					await this._remoteAgentService.endConnection();
+					await this._doStopExtensionHosts();
+					this._remoteAgentService.getConnection()?.dispose();
+				}, {
 					id: 'join.disconnectRemote',
 					label: nls.localize('disconnectRemote', "Disconnect Remote Agent"),
 					order: WillShutdownJoinerOrder.Last // after others have joined that might depend on a remote connection
 				});
+			} else {
+				event.join(this._doStopExtensionHosts(), {
+					id: 'join.stopExtensionHosts',
+					label: nls.localize('stopExtensionHosts', "Stopping Extension Hosts"),
+				});
 			}
 		}));
-
-		this._register(this._lifecycleService.onDidShutdown(() => {
-			// We need to disconnect the management connection before killing the local extension host.
-			// Otherwise, the local extension host might terminate the underlying tunnel before the
-			// management connection has a chance to send its disconnection message.
-			const connection = this._remoteAgentService.getConnection();
-			connection?.dispose();
-
-			this._doStopExtensionHosts();
-		}));
 	}
 
 	protected _getExtensionHostManagers(kind: ExtensionHostKind): IExtensionHostManager[] {
@@ -667,15 +669,15 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 		return this._doStopExtensionHostsWithVeto(reason, auto);
 	}
 
-	protected _doStopExtensionHosts(): void {
+	protected async _doStopExtensionHosts(): Promise<void> {
 		const previouslyActivatedExtensionIds: ExtensionIdentifier[] = [];
 		for (const extensionStatus of this._extensionStatus.values()) {
 			if (extensionStatus.activationStarted) {
 				previouslyActivatedExtensionIds.push(extensionStatus.id);
 			}
 		}
 
-		this._extensionHostManagers.disposeAllInReverse();
+		await this._extensionHostManagers.stopAllInReverse();
 		for (const extensionStatus of this._extensionStatus.values()) {
 			extensionStatus.clearRuntimeStatus();
 		}
@@ -712,7 +714,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 
 		const veto = await handleVetos(vetos, error => this._logService.error(error));
 		if (!veto) {
-			this._doStopExtensionHosts();
+			await this._doStopExtensionHosts();
 		} else {
 			if (!auto) {
 				const vetoReasonsArray = Array.from(vetoReasons);
@@ -803,7 +805,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 			if (signal) {
 				this._onRemoteExtensionHostCrashed(extensionHost, signal);
 			}
-			this._extensionHostManagers.disposeOne(extensionHost);
+			this._extensionHostManagers.stopOne(extensionHost);
 		}
 	}
 
@@ -868,7 +870,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 	}
 
 	public async startExtensionHosts(updates?: { toAdd: IExtension[]; toRemove: string[] }): Promise<void> {
-		this._doStopExtensionHosts();
+		await this._doStopExtensionHosts();
 
 		if (updates) {
 			await this._handleDeltaExtensions(new DeltaExtensionsQueueItem(updates.toAdd, updates.toRemove));
@@ -1182,37 +1184,45 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 	//#endregion
 
 	protected abstract _resolveExtensions(): Promise<ResolvedExtensions>;
-	protected abstract _onExtensionHostExit(code: number): void;
+	protected abstract _onExtensionHostExit(code: number): Promise<void>;
 	protected abstract _resolveAuthority(remoteAuthority: string): Promise<ResolverResult>;
 }
 
 class ExtensionHostCollection extends Disposable {
 
 	private _extensionHostManagers: ExtensionHostManagerData[] = [];
 
-	public override dispose(): void {
-		this.disposeAllInReverse();
+	public override dispose() {
+		for (let i = this._extensionHostManagers.length - 1; i >= 0; i--) {
+			const manager = this._extensionHostManagers[i];
+			manager.extensionHost.disconnect();
+			manager.dispose();
+		}
+		this._extensionHostManagers = [];
 		super.dispose();
 	}
 
 	public add(extensionHostManager: IExtensionHostManager, disposableStore: DisposableStore): void {
 		this._extensionHostManagers.push(new ExtensionHostManagerData(extensionHostManager, disposableStore));
 	}
 
-	public disposeAllInReverse(): void {
+	public async stopAllInReverse(): Promise<void> {
 		// See https://github.com/microsoft/vscode/issues/152204
 		// Dispose extension hosts in reverse creation order because the local extension host
 		// might be critical in sustaining a connection to the remote extension host
 		for (let i = this._extensionHostManagers.length - 1; i >= 0; i--) {
-			this._extensionHostManagers[i].dispose();
+			const manager = this._extensionHostManagers[i];
+			await manager.extensionHost.disconnect();
+			manager.dispose();
 		}
 		this._extensionHostManagers = [];
 	}
 
-	public disposeOne(extensionHostManager: IExtensionHostManager): void {
+	public async stopOne(extensionHostManager: IExtensionHostManager): Promise<void> {
 		const index = this._extensionHostManagers.findIndex(el => el.extensionHost === extensionHostManager);
 		if (index >= 0) {
 			this._extensionHostManagers.splice(index, 1);
+			await extensionHostManager.disconnect();
 			extensionHostManager.dispose();
 		}
 	}

src/vs/workbench/services/extensions/common/extensionHostManager.ts

@@ -161,13 +161,14 @@ export class ExtensionHostManager extends Disposable implements IExtensionHostMa
 		});
 	}
 
+	public async disconnect(): Promise<void> {
+		await this._extensionHost?.disconnect?.();
+	}
+
 	public override dispose(): void {
-		if (this._extensionHost) {
-			this._extensionHost.dispose();
-		}
-		if (this._rpcProtocol) {
-			this._rpcProtocol.dispose();
-		}
+		this._extensionHost?.dispose();
+		this._rpcProtocol?.dispose();
+
 		for (let i = 0, len = this._customers.length; i < len; i++) {
 			const customer = this._customers[i];
 			try {

src/vs/workbench/services/extensions/common/extensionHostManagers.ts

@@ -20,6 +20,7 @@ export interface IExtensionHostManager {
 	readonly friendyName: string;
 	readonly onDidExit: Event<[number, string | null]>;
 	readonly onDidChangeResponsiveState: Event<ResponsiveState>;
+	disconnect(): Promise<void>;
 	dispose(): void;
 	ready(): Promise<void>;
 	representsRunningLocation(runningLocation: ExtensionRunningLocation): boolean;

src/vs/workbench/services/extensions/common/extensions.ts

@@ -127,6 +127,7 @@ export interface IExtensionHost {
 	start(): Promise<IMessagePassingProtocol>;
 	getInspectPort(): { port: number; host: string } | undefined;
 	enableInspectPort(): Promise<boolean>;
+	disconnect?(): Promise<void>;
 	dispose(): void;
 }
 

src/vs/workbench/services/extensions/common/lazyCreateExtensionHostManager.ts

@@ -90,6 +90,9 @@ export class LazyCreateExtensionHostManager extends Disposable implements IExten
 			await this._actual.ready();
 		}
 	}
+	public async disconnect(): Promise<void> {
+		await this._actual?.disconnect();
+	}
 	public representsRunningLocation(runningLocation: ExtensionRunningLocation): boolean {
 		return this._extensionHost.runningLocation.equals(runningLocation);
 	}

src/vs/workbench/services/extensions/common/remoteExtensionHost.ts

@@ -56,6 +56,7 @@ export class RemoteExtensionHost extends Disposable implements IExtensionHost {
 	private _protocol: PersistentProtocol | null;
 	private _hasLostConnection: boolean;
 	private _terminating: boolean;
+	private _hasDisconnected = false;
 	private readonly _isExtensionDevHost: boolean;
 
 	constructor(
@@ -264,22 +265,28 @@ export class RemoteExtensionHost extends Disposable implements IExtensionHost {
 		return Promise.resolve(false);
 	}
 
+	async disconnect() {
+		if (this._protocol && !this._hasDisconnected) {
+			this._protocol.send(createMessageOfType(MessageType.Terminate));
+			this._protocol.sendDisconnect();
+			this._hasDisconnected = true;
+			await this._protocol.drain();
+		}
+	}
+
 	override dispose(): void {
 		super.dispose();
 
 		this._terminating = true;
+		this.disconnect();
 
 		if (this._protocol) {
 			// Send the extension host a request to terminate itself
 			// (graceful termination)
 			// setTimeout(() => {
 			// console.log(`SENDING TERMINATE TO REMOTE EXT HOST!`);
-			const socket = this._protocol.getSocket();
-			this._protocol.send(createMessageOfType(MessageType.Terminate));
-			this._protocol.sendDisconnect();
-			this._protocol.dispose();
+			this._protocol.getSocket().end();
 			// this._protocol.drain();
-			socket.end();
 			this._protocol = null;
 			// }, 1000);
 		}

src/vs/workbench/services/extensions/electron-sandbox/nativeExtensionService.ts

@@ -420,9 +420,9 @@ export class NativeExtensionService extends AbstractExtensionService implements
 		return new ResolvedExtensions(await this._scanAllLocalExtensions(), remoteExtensions, /*hasLocalProcess*/true, /*allowRemoteExtensionsInLocalWebWorker*/false);
 	}
 
-	protected _onExtensionHostExit(code: number): void {
+	protected async _onExtensionHostExit(code: number): Promise<void> {
 		// Dispose everything associated with the extension host
-		this._doStopExtensionHosts();
+		await this._doStopExtensionHosts();
 
 		// Dispose the management connection to avoid reconnecting after the extension host exits
 		const connection = this._remoteAgentService.getConnection();

src/vs/workbench/services/extensions/test/browser/extensionService.test.ts

@@ -198,6 +198,9 @@ suite('ExtensionService', () => {
 			return new class extends mock<IExtensionHostManager>() {
 				override onDidExit = Event.None;
 				override onDidChangeResponsiveState = Event.None;
+				override disconnect() {
+					return Promise.resolve();
+				}
 				override dispose(): void {
 					order.push(`dispose ${extensionHostId}`);
 				}
@@ -212,7 +215,7 @@ suite('ExtensionService', () => {
 		protected _scanSingleExtension(extension: IExtension): Promise<IExtensionDescription | null> {
 			throw new Error('Method not implemented.');
 		}
-		protected _onExtensionHostExit(code: number): void {
+		protected _onExtensionHostExit(code: number): Promise<void> {
 			throw new Error('Method not implemented.');
 		}
 		protected _resolveAuthority(remoteAuthority: string): Promise<ResolverResult> {
@@ -255,7 +258,7 @@ suite('ExtensionService', () => {
 		extService = <MyTestExtensionService>instantiationService.get(IExtensionService);
 	});
 
-	teardown(() => {
+	teardown(async () => {
 		disposables.dispose();
 	});
 

src/vs/workbench/services/lifecycle/common/lifecycleService.ts

@@ -3,13 +3,13 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Barrier } from 'vs/base/common/async';
 import { Emitter } from 'vs/base/common/event';
+import { Barrier } from 'vs/base/common/async';
 import { Disposable } from 'vs/base/common/lifecycle';
-import { mark } from 'vs/base/common/performance';
+import { ILifecycleService, WillShutdownEvent, StartupKind, LifecyclePhase, LifecyclePhaseToString, ShutdownReason, BeforeShutdownErrorEvent, InternalBeforeShutdownEvent } from 'vs/workbench/services/lifecycle/common/lifecycle';
 import { ILogService } from 'vs/platform/log/common/log';
+import { mark } from 'vs/base/common/performance';
 import { IStorageService, StorageScope, StorageTarget, WillSaveStateReason } from 'vs/platform/storage/common/storage';
-import { BeforeShutdownErrorEvent, ILifecycleService, InternalBeforeShutdownEvent, LifecyclePhase, LifecyclePhaseToString, ShutdownReason, StartupKind, WillShutdownEvent } from 'vs/workbench/services/lifecycle/common/lifecycle';
 
 export abstract class AbstractLifecycleService extends Disposable implements ILifecycleService {
 
@@ -44,7 +44,7 @@ export abstract class AbstractLifecycleService extends Disposable implements ILi
 
 	constructor(
 		@ILogService protected readonly logService: ILogService,
-		@IStorageService protected readonly storageService: IStorageService,
+		@IStorageService protected readonly storageService: IStorageService
 	) {
 		super();
 

src/vs/workbench/services/lifecycle/electron-sandbox/lifecycleService.ts

@@ -11,22 +11,19 @@ import { ILogService } from 'vs/platform/log/common/log';
 import { AbstractLifecycleService } from 'vs/workbench/services/lifecycle/common/lifecycleService';
 import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { INativeHostService } from 'vs/platform/native/common/native';
-import { Promises, disposableTimeout, raceCancellation, timeout } from 'vs/base/common/async';
+import { Promises, disposableTimeout, raceCancellation } from 'vs/base/common/async';
 import { toErrorMessage } from 'vs/base/common/errorMessage';
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
-import { IRemoteAgentService } from 'vs/workbench/services/remote/common/remoteAgentService';
 
 export class NativeLifecycleService extends AbstractLifecycleService {
 
 	private static readonly BEFORE_SHUTDOWN_WARNING_DELAY = 5000;
 	private static readonly WILL_SHUTDOWN_WARNING_DELAY = 800;
-	private static readonly MAX_GRACEFUL_REMOTE_DISCONNECT_TIME = 3000;
 
 	constructor(
 		@INativeHostService private readonly nativeHostService: INativeHostService,
 		@IStorageService storageService: IStorageService,
-		@ILogService logService: ILogService,
-		@IRemoteAgentService private readonly remoteAgentService: IRemoteAgentService,
+		@ILogService logService: ILogService
 	) {
 		super(logService, storageService);
 
@@ -69,10 +66,6 @@ export class NativeLifecycleService extends AbstractLifecycleService {
 			// trigger onWillShutdown events and joining
 			await this.handleWillShutdown(reply.reason);
 
-			// now that all services have stored their data, it's safe to terminate
-			// the remote connection gracefully before synchronously saying we're shut down
-			await this.handleRemoteAgentDisconnect();
-
 			// trigger onDidShutdown event now that we know we will quit
 			this._onDidShutdown.fire();
 
@@ -81,19 +74,6 @@ export class NativeLifecycleService extends AbstractLifecycleService {
 		});
 	}
 
-	private async handleRemoteAgentDisconnect(): Promise<void> {
-		const longRunningWarning = disposableTimeout(() => {
-			this.logService.warn(`[lifecycle] the remote agent is taking a long time to disconnect, waiting...`);
-		}, NativeLifecycleService.BEFORE_SHUTDOWN_WARNING_DELAY);
-
-		await Promise.race([
-			this.remoteAgentService.endConnection(),
-			timeout(NativeLifecycleService.MAX_GRACEFUL_REMOTE_DISCONNECT_TIME),
-		]);
-
-		longRunningWarning.dispose();
-	}
-
 	protected async handleBeforeShutdown(reason: ShutdownReason): Promise<boolean> {
 		const logService = this.logService;
 

src/vs/workbench/services/remote/browser/remoteAgentService.ts

@@ -27,7 +27,7 @@ export class RemoteAgentService extends AbstractRemoteAgentService implements IR
 		@IProductService productService: IProductService,
 		@IRemoteAuthorityResolverService remoteAuthorityResolverService: IRemoteAuthorityResolverService,
 		@ISignService signService: ISignService,
-		@ILogService logService: ILogService,
+		@ILogService logService: ILogService
 	) {
 		super(remoteSocketFactoryService, userDataProfileService, environmentService, productService, remoteAuthorityResolverService, signService, logService);
 	}