Merge pull request microsoft#103122 from microsoft/alex/102062

alexdima · web-flow · commit 6f685fdba500 · 2020-07-22T23:52:30.000+02:00
diff --git a/src/vs/base/parts/ipc/common/ipc.net.ts b/src/vs/base/parts/ipc/common/ipc.net.ts
@@ -16,6 +16,7 @@ export interface ISocket extends IDisposable {
 	onEnd(listener: () => void): IDisposable;
 	write(buffer: VSBuffer): void;
 	end(): void;
+	drain(): Promise<void>;
 }
 
 let emptyBuffer: VSBuffer | null = null;
@@ -277,6 +278,11 @@ class ProtocolWriter {
 		this._isDisposed = true;
 	}
 
+	public drain(): Promise<void> {
+		this.flush();
+		return this._socket.drain();
+	}
+
 	public flush(): void {
 		// flush
 		this._writeNow();
@@ -372,6 +378,10 @@ export class Protocol extends Disposable implements IMessagePassingProtocol {
 		this._register(this._socket.onClose(() => this._onClose.fire()));
 	}
 
+	drain(): Promise<void> {
+		return this._socketWriter.drain();
+	}
+
 	getSocket(): ISocket {
 		return this._socket;
 	}
@@ -619,6 +629,10 @@ export class PersistentProtocol implements IMessagePassingProtocol {
 		this._socketDisposables = dispose(this._socketDisposables);
 	}
 
+	drain(): Promise<void> {
+		return this._socketWriter.drain();
+	}
+
 	sendDisconnect(): void {
 		const msg = new ProtocolMessage(ProtocolMessageType.Disconnect, 0, 0, getEmptyBuffer());
 		this._socketWriter.write(msg);
diff --git a/src/vs/base/parts/ipc/common/ipc.ts b/src/vs/base/parts/ipc/common/ipc.ts
@@ -70,6 +70,10 @@ interface IHandler {
 export interface IMessagePassingProtocol {
 	send(buffer: VSBuffer): void;
 	onMessage: Event<VSBuffer>;
+	/**
+	 * Wait for the write buffer (if applicable) to become empty.
+	 */
+	drain?(): Promise<void>;
 }
 
 enum State {
@@ -482,10 +486,7 @@ export class ChannelClient implements IChannelClient, IDisposable {
 				return e(errors.canceled());
 			}
 
-			let uninitializedPromise: CancelablePromise<void> | null = createCancelablePromise(_ => this.whenInitialized());
-			uninitializedPromise.then(() => {
-				uninitializedPromise = null;
-
+			const doRequest = () => {
 				const handler: IHandler = response => {
 					switch (response.type) {
 						case ResponseType.PromiseSuccess:
@@ -510,7 +511,18 @@ export class ChannelClient implements IChannelClient, IDisposable {
 
 				this.handlers.set(id, handler);
 				this.sendRequest(request);
-			});
+			};
+
+			let uninitializedPromise: CancelablePromise<void> | null = null;
+			if (this.state === State.Idle) {
+				doRequest();
+			} else {
+				uninitializedPromise = createCancelablePromise(_ => this.whenInitialized());
+				uninitializedPromise.then(() => {
+					uninitializedPromise = null;
+					doRequest();
+				});
+			}
 
 			const cancel = () => {
 				if (uninitializedPromise) {
diff --git a/src/vs/base/parts/ipc/node/ipc.net.ts b/src/vs/base/parts/ipc/node/ipc.net.ts
@@ -77,6 +77,28 @@ export class NodeSocket implements ISocket {
 	public end(): void {
 		this.socket.end();
 	}
+
+	public drain(): Promise<void> {
+		return new Promise<void>((resolve, reject) => {
+			if (this.socket.bufferSize === 0) {
+				resolve();
+				return;
+			}
+			const finished = () => {
+				this.socket.off('close', finished);
+				this.socket.off('end', finished);
+				this.socket.off('error', finished);
+				this.socket.off('timeout', finished);
+				this.socket.off('drain', finished);
+				resolve();
+			};
+			this.socket.on('close', finished);
+			this.socket.on('end', finished);
+			this.socket.on('error', finished);
+			this.socket.on('timeout', finished);
+			this.socket.on('drain', finished);
+		});
+	}
 }
 
 const enum Constants {
@@ -243,6 +265,10 @@ export class WebSocketNodeSocket extends Disposable implements ISocket {
 			}
 		}
 	}
+
+	public drain(): Promise<void> {
+		return this.socket.drain();
+	}
 }
 
 function unmask(buffer: VSBuffer, mask: number): void {
diff --git a/src/vs/platform/remote/browser/browserSocketFactory.ts b/src/vs/platform/remote/browser/browserSocketFactory.ts
@@ -194,6 +194,9 @@ class BrowserSocket implements ISocket {
 		this.socket.close();
 	}
 
+	public drain(): Promise<void> {
+		return Promise.resolve();
+	}
 }
 
 
diff --git a/src/vs/workbench/api/browser/mainThreadExtensionService.ts b/src/vs/workbench/api/browser/mainThreadExtensionService.ts
@@ -128,7 +128,7 @@ export class MainThreadExtensionService implements MainThreadExtensionServiceSha
 		}
 	}
 
-	$onExtensionHostExit(code: number): void {
+	async $onExtensionHostExit(code: number): Promise<void> {
 		this._extensionService._onExtensionHostExit(code);
 	}
 }
diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts
@@ -795,7 +795,7 @@ export interface MainThreadExtensionServiceShape extends IDisposable {
 	$onDidActivateExtension(extensionId: ExtensionIdentifier, codeLoadingTime: number, activateCallTime: number, activateResolvedTime: number, activationReason: ExtensionActivationReason): void;
 	$onExtensionActivationError(extensionId: ExtensionIdentifier, error: ExtensionActivationError): Promise<void>;
 	$onExtensionRuntimeError(extensionId: ExtensionIdentifier, error: SerializedError): void;
-	$onExtensionHostExit(code: number): void;
+	$onExtensionHostExit(code: number): Promise<void>;
 }
 
 export interface SCMProviderFeatures {
diff --git a/src/vs/workbench/api/common/extHostExtensionService.ts b/src/vs/workbench/api/common/extHostExtensionService.ts
@@ -557,7 +557,7 @@ export abstract class AbstractExtHostExtensionService extends Disposable impleme
 					}
 
 					// after tests have run, we shutdown the host
-					this._gracefulExit(error || (typeof failures === 'number' && failures > 0) ? 1 /* ERROR */ : 0 /* OK */);
+					this._testRunnerExit(error || (typeof failures === 'number' && failures > 0) ? 1 /* ERROR */ : 0 /* OK */);
 				};
 
 				const runResult = testRunner!.run(extensionTestsPath, oldTestRunnerCallback);
@@ -567,36 +567,32 @@ export abstract class AbstractExtHostExtensionService extends Disposable impleme
 					runResult
 						.then(() => {
 							c();
-							this._gracefulExit(0);
+							this._testRunnerExit(0);
 						})
 						.catch((err: Error) => {
 							e(err.toString());
-							this._gracefulExit(1);
+							this._testRunnerExit(1);
 						});
 				}
 			});
 		}
 
 		// Otherwise make sure to shutdown anyway even in case of an error
 		else {
-			this._gracefulExit(1 /* ERROR */);
+			this._testRunnerExit(1 /* ERROR */);
 		}
 
 		return Promise.reject(new Error(requireError ? requireError.toString() : nls.localize('extensionTestError', "Path {0} does not point to a valid extension test runner.", extensionTestsPath)));
 	}
 
-	private _gracefulExit(code: number): void {
-		// to give the PH process a chance to flush any outstanding console
-		// messages to the main process, we delay the exit() by some time
-		setTimeout(() => {
-			// If extension tests are running, give the exit code to the renderer
-			if (this._initData.remote.isRemote && !!this._initData.environment.extensionTestsLocationURI) {
-				this._mainThreadExtensionsProxy.$onExtensionHostExit(code);
-				return;
-			}
-
+	private _testRunnerExit(code: number): void {
+		// wait at most 5000ms for the renderer to confirm our exit request and for the renderer socket to drain
+		// (this is to ensure all outstanding messages reach the renderer)
+		const exitPromise = this._mainThreadExtensionsProxy.$onExtensionHostExit(code);
+		const drainPromise = this._extHostContext.drain();
+		Promise.race([Promise.all([exitPromise, drainPromise]), timeout(5000)]).then(() => {
 			this._hostUtils.exit(code);
-		}, 500);
+		});
 	}
 
 	private _startExtensionHost(): Promise<void> {
diff --git a/src/vs/workbench/api/common/extHostRpcService.ts b/src/vs/workbench/api/common/extHostRpcService.ts
@@ -18,12 +18,12 @@ export class ExtHostRpcService implements IExtHostRpcService {
 	readonly getProxy: <T>(identifier: ProxyIdentifier<T>) => T;
 	readonly set: <T, R extends T> (identifier: ProxyIdentifier<T>, instance: R) => R;
 	readonly assertRegistered: (identifiers: ProxyIdentifier<any>[]) => void;
+	readonly drain: () => Promise<void>;
 
 	constructor(rpcProtocol: IRPCProtocol) {
 		this.getProxy = rpcProtocol.getProxy.bind(rpcProtocol);
 		this.set = rpcProtocol.set.bind(rpcProtocol);
 		this.assertRegistered = rpcProtocol.assertRegistered.bind(rpcProtocol);
-
+		this.drain = rpcProtocol.drain.bind(rpcProtocol);
 	}
-
 }
diff --git a/src/vs/workbench/services/extensions/common/extensionHostManager.ts b/src/vs/workbench/services/extensions/common/extensionHostManager.ts
@@ -184,6 +184,7 @@ export class ExtensionHostManager extends Disposable {
 			getProxy: <T>(identifier: ProxyIdentifier<T>): T => this._rpcProtocol!.getProxy(identifier),
 			set: <T, R extends T>(identifier: ProxyIdentifier<T>, instance: R): R => this._rpcProtocol!.set(identifier, instance),
 			assertRegistered: (identifiers: ProxyIdentifier<any>[]): void => this._rpcProtocol!.assertRegistered(identifiers),
+			drain: (): Promise<void> => this._rpcProtocol!.drain(),
 		};
 
 		// Named customers
diff --git a/src/vs/workbench/services/extensions/common/proxyIdentifier.ts b/src/vs/workbench/services/extensions/common/proxyIdentifier.ts
@@ -18,6 +18,11 @@ export interface IRPCProtocol {
 	 * Assert these identifiers are already registered via `.set`.
 	 */
 	assertRegistered(identifiers: ProxyIdentifier<any>[]): void;
+
+	/**
+	 * Wait for the write buffer (if applicable) to become empty.
+	 */
+	drain(): Promise<void>;
 }
 
 export class ProxyIdentifier<T> {
diff --git a/src/vs/workbench/services/extensions/common/rpcProtocol.ts b/src/vs/workbench/services/extensions/common/rpcProtocol.ts
@@ -115,6 +115,13 @@ export class RPCProtocol extends Disposable implements IRPCProtocol {
 		});
 	}
 
+	public drain(): Promise<void> {
+		if (typeof this._protocol.drain === 'function') {
+			return this._protocol.drain();
+		}
+		return Promise.resolve();
+	}
+
 	private _onWillSendRequest(req: number): void {
 		if (this._unacknowledgedCount === 0) {
 			// Since this is the first request we are sending in a while,
diff --git a/src/vs/workbench/services/extensions/node/extensionHostProcessSetup.ts b/src/vs/workbench/services/extensions/node/extensionHostProcessSetup.ts
@@ -96,10 +96,10 @@ let onTerminate = function () {
 	nativeExit();
 };
 
-function _createExtHostProtocol(): Promise<IMessagePassingProtocol> {
+function _createExtHostProtocol(): Promise<PersistentProtocol> {
 	if (process.env.VSCODE_EXTHOST_WILL_SEND_SOCKET) {
 
-		return new Promise<IMessagePassingProtocol>((resolve, reject) => {
+		return new Promise<PersistentProtocol>((resolve, reject) => {
 
 			let protocol: PersistentProtocol | null = null;
 
@@ -163,7 +163,7 @@ function _createExtHostProtocol(): Promise<IMessagePassingProtocol> {
 
 		const pipeName = process.env.VSCODE_IPC_HOOK_EXTHOST!;
 
-		return new Promise<IMessagePassingProtocol>((resolve, reject) => {
+		return new Promise<PersistentProtocol>((resolve, reject) => {
 
 			const socket = net.createConnection(pipeName, () => {
 				socket.removeListener('error', reject);
@@ -203,6 +203,10 @@ async function createExtHostProtocol(): Promise<IMessagePassingProtocol> {
 				protocol.send(msg);
 			}
 		}
+
+		drain(): Promise<void> {
+			return protocol.drain();
+		}
 	};
 }
 
diff --git a/src/vs/workbench/test/browser/api/extHostDiagnostics.test.ts b/src/vs/workbench/test/browser/api/extHostDiagnostics.test.ts
@@ -389,6 +389,9 @@ suite('ExtHostDiagnostics', () => {
 			assertRegistered(): void {
 
 			}
+			drain() {
+				return undefined!;
+			}
 		}, new NullLogService());
 
 		let collection1 = diags.createDiagnosticCollection(nullExtensionDescription.identifier, 'foo');
@@ -438,6 +441,9 @@ suite('ExtHostDiagnostics', () => {
 			assertRegistered(): void {
 
 			}
+			drain() {
+				return undefined!;
+			}
 		}, new NullLogService());
 
 
diff --git a/src/vs/workbench/test/browser/api/extHostFileSystemEventService.test.ts b/src/vs/workbench/test/browser/api/extHostFileSystemEventService.test.ts
@@ -15,7 +15,8 @@ suite('ExtHostFileSystemEventService', () => {
 		const protocol: IMainContext = {
 			getProxy: () => { return undefined!; },
 			set: undefined!,
-			assertRegistered: undefined!
+			assertRegistered: undefined!,
+			drain: undefined!
 		};
 
 		const watcher1 = new ExtHostFileSystemEventService(protocol, new NullLogService(), undefined!).createFileSystemWatcher('**/somethingInteresting', false, false, false);
diff --git a/src/vs/workbench/test/browser/api/extHostWorkspace.test.ts b/src/vs/workbench/test/browser/api/extHostWorkspace.test.ts
@@ -298,7 +298,8 @@ suite('ExtHostWorkspace', function () {
 		const protocol: IMainContext = {
 			getProxy: () => { return undefined!; },
 			set: () => { return undefined!; },
-			assertRegistered: () => { }
+			assertRegistered: () => { },
+			drain: () => { return undefined!; },
 		};
 
 		const ws = createExtHostWorkspace(protocol, { id: 'foo', name: 'Test', folders: [] }, new NullLogService());
diff --git a/src/vs/workbench/test/browser/api/mainThreadDiagnostics.test.ts b/src/vs/workbench/test/browser/api/mainThreadDiagnostics.test.ts
@@ -32,6 +32,7 @@ suite('MainThreadDiagnostics', function () {
 						$acceptMarkersChange() { }
 					};
 				}
+				drain(): any { return null; }
 			},
 			markerService,
 			new class extends mock<IUriIdentityService>() {
diff --git a/src/vs/workbench/test/browser/api/mainThreadTreeViews.test.ts b/src/vs/workbench/test/browser/api/mainThreadTreeViews.test.ts
@@ -66,6 +66,7 @@ suite('MainThreadHostTreeView', function () {
 				getProxy(): any {
 					return extHostTreeViewsShape;
 				}
+				drain(): any { return null; }
 			}, new TestViewsService(), new TestNotificationService(), testExtensionService, new NullLogService());
 		mainThreadTreeViews.$registerTreeViewDataProvider(testTreeViewId, { showCollapseAll: false, canSelectMany: false });
 		await testExtensionService.whenInstalledExtensionsRegistered();
diff --git a/src/vs/workbench/test/browser/api/testRPCProtocol.ts b/src/vs/workbench/test/browser/api/testRPCProtocol.ts
@@ -19,7 +19,8 @@ export function SingleProxyRPCProtocol(thing: any): IExtHostContext & IExtHostRp
 		set<T, R extends T>(identifier: ProxyIdentifier<T>, value: R): R {
 			return value;
 		},
-		assertRegistered: undefined!
+		assertRegistered: undefined!,
+		drain: undefined!
 	};
 }
 
@@ -40,6 +41,10 @@ export class TestRPCProtocol implements IExtHostContext, IExtHostRpcService {
 		this._proxies = Object.create(null);
 	}
 
+	drain(): Promise<void> {
+		return Promise.resolve();
+	}
+
 	private get _callCount(): number {
 		return this._callCountValue;
 	}

Original file line number	Diff line number	Diff line change
`@@ -194,6 +194,9 @@ class BrowserSocket implements ISocket {`
`194`	`194`	`this.socket.close();`
`195`	`195`	`}`
`196`	`196`
	`197`	`+ public drain(): Promise<void> {`
	`198`	`+ return Promise.resolve();`
	`199`	`+ }`
`197`	`200`	`}`
`198`	`201`
`199`	`202`
Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ export class MainThreadExtensionService implements MainThreadExtensionServiceSha`
`128`	`128`	`}`
`129`	`129`	`}`
`130`	`130`
`131`		`- $onExtensionHostExit(code: number): void {`
	`131`	`+ async $onExtensionHostExit(code: number): Promise<void> {`
`132`	`132`	`this._extensionService._onExtensionHostExit(code);`
`133`	`133`	`}`
`134`	`134`	`}`
Original file line number	Diff line number	Diff line change
`@@ -795,7 +795,7 @@ export interface MainThreadExtensionServiceShape extends IDisposable {`
`795`	`795`	`$onDidActivateExtension(extensionId: ExtensionIdentifier, codeLoadingTime: number, activateCallTime: number, activateResolvedTime: number, activationReason: ExtensionActivationReason): void;`
`796`	`796`	`$onExtensionActivationError(extensionId: ExtensionIdentifier, error: ExtensionActivationError): Promise<void>;`
`797`	`797`	`$onExtensionRuntimeError(extensionId: ExtensionIdentifier, error: SerializedError): void;`
`798`		`- $onExtensionHostExit(code: number): void;`
	`798`	`+ $onExtensionHostExit(code: number): Promise<void>;`
`799`	`799`	`}`
`800`	`800`
`801`	`801`	`export interface SCMProviderFeatures {`
Original file line number	Diff line number	Diff line change
`@@ -18,12 +18,12 @@ export class ExtHostRpcService implements IExtHostRpcService {`
`18`	`18`	`readonly getProxy: <T>(identifier: ProxyIdentifier<T>) => T;`
`19`	`19`	`readonly set: <T, R extends T> (identifier: ProxyIdentifier<T>, instance: R) => R;`
`20`	`20`	`readonly assertRegistered: (identifiers: ProxyIdentifier<any>[]) => void;`
	`21`	`+ readonly drain: () => Promise<void>;`
`21`	`22`
`22`	`23`	`constructor(rpcProtocol: IRPCProtocol) {`
`23`	`24`	`this.getProxy = rpcProtocol.getProxy.bind(rpcProtocol);`
`24`	`25`	`this.set = rpcProtocol.set.bind(rpcProtocol);`
`25`	`26`	`this.assertRegistered = rpcProtocol.assertRegistered.bind(rpcProtocol);`
`26`		`-`
	`27`	`+ this.drain = rpcProtocol.drain.bind(rpcProtocol);`
`27`	`28`	`}`
`28`		`-`
`29`	`29`	`}`