Skip to content

Commit a5d3d4f

Browse files
alexdimaalexdima
committed
Fixes microsoft#102062: Wait for the renderer socket to drain before exiting the extension host process when running tests
1 parent 09904a6 commit a5d3d4f

14 files changed

Lines changed: 103 additions & 29 deletions

File tree

‎src/vs/base/parts/ipc/common/ipc.net.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ export interface ISocket extends IDisposable {
1616
onEnd(listener: () => void): IDisposable;
1717
write(buffer: VSBuffer): void;
1818
end(): void;
19+
drain(): Promise<void>;
1920
}
2021

2122
let emptyBuffer: VSBuffer | null = null;
@@ -277,6 +278,11 @@ class ProtocolWriter {
277278
this._isDisposed = true;
278279
}
279280

281+
public drain(): Promise<void> {
282+
this.flush();
283+
return this._socket.drain();
284+
}
285+
280286
public flush(): void {
281287
// flush
282288
this._writeNow();
@@ -372,6 +378,10 @@ export class Protocol extends Disposable implements IMessagePassingProtocol {
372378
this._register(this._socket.onClose(() => this._onClose.fire()));
373379
}
374380

381+
drain(): Promise<void> {
382+
return this._socketWriter.drain();
383+
}
384+
375385
getSocket(): ISocket {
376386
return this._socket;
377387
}
@@ -619,6 +629,10 @@ export class PersistentProtocol implements IMessagePassingProtocol {
619629
this._socketDisposables = dispose(this._socketDisposables);
620630
}
621631

632+
drain(): Promise<void> {
633+
return this._socketWriter.drain();
634+
}
635+
622636
sendDisconnect(): void {
623637
const msg = new ProtocolMessage(ProtocolMessageType.Disconnect, 0, 0, getEmptyBuffer());
624638
this._socketWriter.write(msg);

‎src/vs/base/parts/ipc/common/ipc.ts‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ interface IHandler {
7070
export interface IMessagePassingProtocol {
7171
send(buffer: VSBuffer): void;
7272
onMessage: Event<VSBuffer>;
73+
/**
74+
* Wait for the write buffer (if applicable) to become empty.
75+
*/
76+
drain?(): Promise<void>;
7377
}
7478

7579
enum State {
@@ -482,10 +486,7 @@ export class ChannelClient implements IChannelClient, IDisposable {
482486
return e(errors.canceled());
483487
}
484488

485-
let uninitializedPromise: CancelablePromise<void> | null = createCancelablePromise(_ => this.whenInitialized());
486-
uninitializedPromise.then(() => {
487-
uninitializedPromise = null;
488-
489+
const doRequest = () => {
489490
const handler: IHandler = response => {
490491
switch (response.type) {
491492
case ResponseType.PromiseSuccess:
@@ -510,7 +511,18 @@ export class ChannelClient implements IChannelClient, IDisposable {
510511

511512
this.handlers.set(id, handler);
512513
this.sendRequest(request);
513-
});
514+
};
515+
516+
let uninitializedPromise: CancelablePromise<void> | null = null;
517+
if (this.state === State.Idle) {
518+
doRequest();
519+
} else {
520+
uninitializedPromise = createCancelablePromise(_ => this.whenInitialized());
521+
uninitializedPromise.then(() => {
522+
uninitializedPromise = null;
523+
doRequest();
524+
});
525+
}
514526

515527
const cancel = () => {
516528
if (uninitializedPromise) {

‎src/vs/base/parts/ipc/node/ipc.net.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,28 @@ export class NodeSocket implements ISocket {
7777
public end(): void {
7878
this.socket.end();
7979
}
80+
81+
public drain(): Promise<void> {
82+
return new Promise<void>((resolve, reject) => {
83+
if (this.socket.bufferSize === 0) {
84+
resolve();
85+
return;
86+
}
87+
const finished = () => {
88+
this.socket.off('close', finished);
89+
this.socket.off('end', finished);
90+
this.socket.off('error', finished);
91+
this.socket.off('timeout', finished);
92+
this.socket.off('drain', finished);
93+
resolve();
94+
};
95+
this.socket.on('close', finished);
96+
this.socket.on('end', finished);
97+
this.socket.on('error', finished);
98+
this.socket.on('timeout', finished);
99+
this.socket.on('drain', finished);
100+
});
101+
}
80102
}
81103

82104
const enum Constants {
@@ -243,6 +265,10 @@ export class WebSocketNodeSocket extends Disposable implements ISocket {
243265
}
244266
}
245267
}
268+
269+
public drain(): Promise<void> {
270+
return this.socket.drain();
271+
}
246272
}
247273

248274
function unmask(buffer: VSBuffer, mask: number): void {

‎src/vs/platform/remote/browser/browserSocketFactory.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,9 @@ class BrowserSocket implements ISocket {
194194
this.socket.close();
195195
}
196196

197+
public drain(): Promise<void> {
198+
return Promise.resolve();
199+
}
197200
}
198201

199202

‎src/vs/workbench/api/browser/mainThreadExtensionService.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ export class MainThreadExtensionService implements MainThreadExtensionServiceSha
128128
}
129129
}
130130

131-
$onExtensionHostExit(code: number): void {
131+
async $onExtensionHostExit(code: number): Promise<void> {
132132
this._extensionService._onExtensionHostExit(code);
133133
}
134134
}

‎src/vs/workbench/api/common/extHost.protocol.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -795,7 +795,7 @@ export interface MainThreadExtensionServiceShape extends IDisposable {
795795
$onDidActivateExtension(extensionId: ExtensionIdentifier, codeLoadingTime: number, activateCallTime: number, activateResolvedTime: number, activationReason: ExtensionActivationReason): void;
796796
$onExtensionActivationError(extensionId: ExtensionIdentifier, error: ExtensionActivationError): Promise<void>;
797797
$onExtensionRuntimeError(extensionId: ExtensionIdentifier, error: SerializedError): void;
798-
$onExtensionHostExit(code: number): void;
798+
$onExtensionHostExit(code: number): Promise<void>;
799799
}
800800

801801
export interface SCMProviderFeatures {

‎src/vs/workbench/api/common/extHostExtensionService.ts‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -557,7 +557,7 @@ export abstract class AbstractExtHostExtensionService extends Disposable impleme
557557
}
558558

559559
// after tests have run, we shutdown the host
560-
this._gracefulExit(error || (typeof failures === 'number' && failures > 0) ? 1 /* ERROR */ : 0 /* OK */);
560+
this._testRunnerExit(error || (typeof failures === 'number' && failures > 0) ? 1 /* ERROR */ : 0 /* OK */);
561561
};
562562

563563
const runResult = testRunner!.run(extensionTestsPath, oldTestRunnerCallback);
@@ -567,36 +567,32 @@ export abstract class AbstractExtHostExtensionService extends Disposable impleme
567567
runResult
568568
.then(() => {
569569
c();
570-
this._gracefulExit(0);
570+
this._testRunnerExit(0);
571571
})
572572
.catch((err: Error) => {
573573
e(err.toString());
574-
this._gracefulExit(1);
574+
this._testRunnerExit(1);
575575
});
576576
}
577577
});
578578
}
579579

580580
// Otherwise make sure to shutdown anyway even in case of an error
581581
else {
582-
this._gracefulExit(1 /* ERROR */);
582+
this._testRunnerExit(1 /* ERROR */);
583583
}
584584

585585
return Promise.reject(new Error(requireError ? requireError.toString() : nls.localize('extensionTestError', "Path {0} does not point to a valid extension test runner.", extensionTestsPath)));
586586
}
587587

588-
private _gracefulExit(code: number): void {
589-
// to give the PH process a chance to flush any outstanding console
590-
// messages to the main process, we delay the exit() by some time
591-
setTimeout(() => {
592-
// If extension tests are running, give the exit code to the renderer
593-
if (this._initData.remote.isRemote && !!this._initData.environment.extensionTestsLocationURI) {
594-
this._mainThreadExtensionsProxy.$onExtensionHostExit(code);
595-
return;
596-
}
597-
588+
private _testRunnerExit(code: number): void {
589+
// wait at most 5000ms for the renderer to confirm our exit request and for the renderer socket to drain
590+
// (this is to ensure all outstanding messages reach the renderer)
591+
const exitPromise = this._mainThreadExtensionsProxy.$onExtensionHostExit(code);
592+
const drainPromise = this._extHostContext.drain();
593+
Promise.race([Promise.all([exitPromise, drainPromise]), timeout(5000)]).then(() => {
598594
this._hostUtils.exit(code);
599-
}, 500);
595+
});
600596
}
601597

602598
private _startExtensionHost(): Promise<void> {

‎src/vs/workbench/api/common/extHostRpcService.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,12 @@ export class ExtHostRpcService implements IExtHostRpcService {
1818
readonly getProxy: <T>(identifier: ProxyIdentifier<T>) => T;
1919
readonly set: <T, R extends T> (identifier: ProxyIdentifier<T>, instance: R) => R;
2020
readonly assertRegistered: (identifiers: ProxyIdentifier<any>[]) => void;
21+
readonly drain: () => Promise<void>;
2122

2223
constructor(rpcProtocol: IRPCProtocol) {
2324
this.getProxy = rpcProtocol.getProxy.bind(rpcProtocol);
2425
this.set = rpcProtocol.set.bind(rpcProtocol);
2526
this.assertRegistered = rpcProtocol.assertRegistered.bind(rpcProtocol);
26-
27+
this.drain = rpcProtocol.drain.bind(rpcProtocol);
2728
}
28-
2929
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@ export class ExtensionHostManager extends Disposable {
184184
getProxy: <T>(identifier: ProxyIdentifier<T>): T => this._rpcProtocol!.getProxy(identifier),
185185
set: <T, R extends T>(identifier: ProxyIdentifier<T>, instance: R): R => this._rpcProtocol!.set(identifier, instance),
186186
assertRegistered: (identifiers: ProxyIdentifier<any>[]): void => this._rpcProtocol!.assertRegistered(identifiers),
187+
drain: (): Promise<void> => this._rpcProtocol!.drain(),
187188
};
188189

189190
// Named customers

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ export interface IRPCProtocol {
1818
* Assert these identifiers are already registered via `.set`.
1919
*/
2020
assertRegistered(identifiers: ProxyIdentifier<any>[]): void;
21+
22+
/**
23+
* Wait for the write buffer (if applicable) to become empty.
24+
*/
25+
drain(): Promise<void>;
2126
}
2227

2328
export class ProxyIdentifier<T> {

0 commit comments

Comments
 (0)