extension host - allow to veto stopping (#179224) (#180513)

src/vs/workbench/contrib/extensions/browser/extensionEnablementWorkspaceTrustTransitionParticipant.ts

@@ -38,7 +38,7 @@ export class ExtensionEnablementWorkspaceTrustTransitionParticipant extends Disp
  if (environmentService.remoteAuthority) {
  hostService.reload();
  } else {
- extensionService.stopExtensionHosts();
+ extensionService.stopExtensionHosts(true); // TODO@lszomoru adopt support for extension host to veto stopping
  await extensionEnablementService.updateExtensionsEnablementsWhenWorkspaceTrustChanges();
  extensionService.startExtensionHosts();
  }

src/vs/workbench/contrib/relauncher/browser/relauncher.contribution.ts

@@ -196,15 +196,18 @@ export class WorkspaceChangeExtHostRelauncher extends Disposable implements IWor
  ) {
  super();
 
- this.extensionHostRestarter = this._register(new RunOnceScheduler(() => {
+ this.extensionHostRestarter = this._register(new RunOnceScheduler(async () => {
  if (!!environmentService.extensionTestsLocationURI) {
  return; // no restart when in tests: see https://github.com/microsoft/vscode/issues/66936
  }
 
  if (environmentService.remoteAuthority) {
  hostService.reload(); // TODO@aeschli, workaround
  } else if (isNative) {
- extensionService.restartExtensionHost();
+ const stopped = await extensionService.stopExtensionHosts();
+ if (stopped) {
+ extensionService.startExtensionHosts();
+ }
  }
  }, 10));
 

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

@@ -141,7 +141,7 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 
  protected _onExtensionHostExit(code: number): void {
  // Dispose everything associated with the extension host
- this.stopExtensionHosts();
+ this._doStopExtensionHosts();
 
  const automatedWindow = window as unknown as IAutomatedWindow;
  if (typeof automatedWindow.codeAutomationExit === 'function') {

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

@@ -17,6 +17,7 @@ import { ImplicitActivationEvents } from 'vs/platform/extensionManagement/common
 import { ExtensionIdentifier, ExtensionIdentifierMap, IExtension, IExtensionContributions, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { IFileService } from 'vs/platform/files/common/files';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { handleVetos } from 'vs/platform/lifecycle/common/lifecycle';
 import { ILogService } from 'vs/platform/log/common/log';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
 import { IProductService } from 'vs/platform/product/common/productService';
@@ -32,7 +33,7 @@ import { IExtensionHostManager, createExtensionHostManager } from 'vs/workbench/
 import { IExtensionManifestPropertiesService } from 'vs/workbench/services/extensions/common/extensionManifestPropertiesService';
 import { ExtensionRunningLocation, LocalProcessRunningLocation, LocalWebWorkerRunningLocation, RemoteRunningLocation } from 'vs/workbench/services/extensions/common/extensionRunningLocation';
 import { ExtensionRunningLocationTracker, filterExtensionIdentifiers } from 'vs/workbench/services/extensions/common/extensionRunningLocationTracker';
-import { ActivationKind, ActivationTimes, ExtensionActivationReason, ExtensionHostStartup, ExtensionPointContribution, IExtensionHost, IExtensionService, IExtensionsStatus, IInternalExtensionService, IMessage, IResponsiveStateChangeEvent, IWillActivateEvent, toExtension } from 'vs/workbench/services/extensions/common/extensions';
+import { ActivationKind, ActivationTimes, ExtensionActivationReason, ExtensionHostStartup, ExtensionPointContribution, IExtensionHost, IExtensionService, IExtensionsStatus, IInternalExtensionService, IMessage, IResponsiveStateChangeEvent, IWillActivateEvent, WillStopExtensionHostsEvent, toExtension } from 'vs/workbench/services/extensions/common/extensions';
 import { ExtensionsProposedApi } from 'vs/workbench/services/extensions/common/extensionsProposedApi';
 import { ExtensionMessageCollector, ExtensionPoint, ExtensionsRegistry, IExtensionPoint, IExtensionPointUser } from 'vs/workbench/services/extensions/common/extensionsRegistry';
 import { ResponsiveState } from 'vs/workbench/services/extensions/common/rpcProtocol';
@@ -62,6 +63,9 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  private readonly _onDidChangeResponsiveChange = this._register(new Emitter<IResponsiveStateChangeEvent>());
  public readonly onDidChangeResponsiveChange = this._onDidChangeResponsiveChange.event;
 
+ private readonly _onWillStop = this._register(new Emitter<WillStopExtensionHostsEvent>());
+ public readonly onWillStop = this._onWillStop.event;
+
  private readonly _activationEventReader = new ImplicitActivationAwareReader();
  private readonly _registry = new LockableExtensionDescriptionRegistry(this._activationEventReader);
  private readonly _installedExtensionsReady = new Barrier();
@@ -160,7 +164,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  const connection = this._remoteAgentService.getConnection();
  connection?.dispose();
 
- this.stopExtensionHosts();
+ this._doStopExtensionHosts();
  }));
  }
 
@@ -519,7 +523,17 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 
  //#region Stopping / Starting / Restarting
 
- public stopExtensionHosts(): void {
+ public stopExtensionHosts(): Promise<boolean>;
+ public stopExtensionHosts(force: true): void;
+ public stopExtensionHosts(force?: boolean): void | Promise<boolean> {
+ if (force) {
+ return this._doStopExtensionHosts();
+ }
+
+ return this._doStopExtensionHostsWithVeto();
+ }
+
+ protected _doStopExtensionHosts(): void {
  const previouslyActivatedExtensionIds: ExtensionIdentifier[] = [];
  for (const extensionStatus of this._extensionStatus.values()) {
  if (extensionStatus.activationStarted) {
@@ -543,6 +557,25 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  }
  }
 
+ private async _doStopExtensionHostsWithVeto(): Promise<boolean> {
+ const vetos: (boolean | Promise<boolean>)[] = [];
+
+ this._onWillStop.fire({
+ veto(value) {
+ vetos.push(value);
+ }
+ });
+
+ const veto = await handleVetos(vetos, error => this._logService.error(error));
+ if (!veto) {
+ this._doStopExtensionHosts();
+ } else {
+ this._logService.warn('Extension Host stop request was vetoed');
+ }
+
+ return !veto;
+ }
+
  private _startExtensionHostsIfNecessary(isInitialStart: boolean, initialActivationEvents: string[]): void {
  const locations: ExtensionRunningLocation[] = [];
  for (let affinity = 0; affinity <= this._runningLocations.maxLocalProcessAffinity; affinity++) {
@@ -603,7 +636,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  protected _onExtensionHostCrashed(extensionHost: IExtensionHostManager, code: number, signal: string | null): void {
  console.error(`Extension host (${extensionHostKindToString(extensionHost.kind)}) terminated unexpectedly. Code: ${code}, Signal: ${signal}`);
  if (extensionHost.kind === ExtensionHostKind.LocalProcess) {
- this.stopExtensionHosts();
+ this._doStopExtensionHosts();
  } else if (extensionHost.kind === ExtensionHostKind.Remote) {
  if (signal) {
  this._onRemoteExtensionHostCrashed(extensionHost, signal);
@@ -679,7 +712,7 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  }
 
  public async startExtensionHosts(): Promise<void> {
- this.stopExtensionHosts();
+ this._doStopExtensionHosts();
 
  const lock = await this._registry.acquireLock('startExtensionHosts');
  try {
@@ -692,11 +725,6 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
  }
  }
 
- public async restartExtensionHost(): Promise<void> {
- this.stopExtensionHosts();
- await this.startExtensionHosts();
- }
-
  //#endregion
 
  //#region IExtensionService

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

@@ -323,6 +323,18 @@ export const enum ActivationKind {
  Immediate = 1
 }
 
+export interface WillStopExtensionHostsEvent {
+
+ /**
+ * Allows to veto the stopping of extension hosts. The veto can be a long running
+ * operation.
+ *
+ * @param id to identify the veto operation in case it takes very long or never
+ * completes.
+ */
+ veto(value: boolean | Promise<boolean>, id: string): void;
+}
+
 export interface IExtensionService {
  readonly _serviceBrand: undefined;
 
@@ -366,6 +378,12 @@ export interface IExtensionService {
  */
  onDidChangeResponsiveChange: Event<IResponsiveStateChangeEvent>;
 
+ /**
+ * Fired before stop of extension hosts happens. Allows listeners to veto against the
+ * stop to prevent it from happening.
+ */
+ onWillStop: Event<WillStopExtensionHostsEvent>;
+
  /**
  * Send an activation event and activate interested extensions.
  *
@@ -426,13 +444,16 @@ export interface IExtensionService {
 
  /**
  * Stops the extension hosts.
+ *
+ * @returns a promise that resolves to `true` if the extension hosts were stopped, `false`
+ * if the operation was vetoed by listeners of the `onWillStop` event.
  */
- stopExtensionHosts(): void;
+ stopExtensionHosts(): Promise<boolean>;
 
  /**
- * Restarts the extension host.
+ * @deprecated Use `stopExtensionHosts()` instead.
  */
- restartExtensionHost(): Promise<void>;
+ stopExtensionHosts(force: true): void;
 
  /**
  * Starts the extension hosts.
@@ -492,6 +513,7 @@ export class NullExtensionService implements IExtensionService {
  onDidChangeExtensions = Event.None;
  onWillActivateByEvent: Event<IWillActivateEvent> = Event.None;
  onDidChangeResponsiveChange: Event<IResponsiveStateChangeEvent> = Event.None;
+ onWillStop: Event<WillStopExtensionHostsEvent> = Event.None;
  readonly extensions = [];
  activateByEvent(_activationEvent: string): Promise<void> { return Promise.resolve(undefined); }
  activationEventIsDone(_activationEvent: string): boolean { return false; }
@@ -500,8 +522,7 @@ export class NullExtensionService implements IExtensionService {
  readExtensionPointContributions<T>(_extPoint: IExtensionPoint<T>): Promise<ExtensionPointContribution<T>[]> { return Promise.resolve(Object.create(null)); }
  getExtensionsStatus(): { [id: string]: IExtensionsStatus } { return Object.create(null); }
  getInspectPorts(_extensionHostKind: ExtensionHostKind, _tryEnableInspector: boolean): Promise<number[]> { return Promise.resolve([]); }
- stopExtensionHosts(): void { }
- async restartExtensionHost(): Promise<void> { }
+ stopExtensionHosts(): any { }
  async startExtensionHosts(): Promise<void> { }
  async setRemoteEnvironment(_env: { [key: string]: string | null }): Promise<void> { }
  canAddExtension(): boolean { return false; }

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

@@ -506,7 +506,7 @@ export class NativeExtensionService extends AbstractExtensionService implements
 
  protected _onExtensionHostExit(code: number): void {
  // Dispose everything associated with the extension host
- this.stopExtensionHosts();
+ this._doStopExtensionHosts();
 
  // Dispose the management connection to avoid reconnecting after the extension host exits
  const connection = this._remoteAgentService.getConnection();
@@ -799,8 +799,13 @@ class RestartExtensionHostAction extends Action2 {
  });
  }
 
- run(accessor: ServicesAccessor): void {
- accessor.get(IExtensionService).restartExtensionHost();
+ async run(accessor: ServicesAccessor): Promise<void> {
+ const extensionService = accessor.get(IExtensionService);
+
+ const stopped = await extensionService.stopExtensionHosts();
+ if (stopped) {
+ extensionService.startExtensionHosts();
+ }
  }
 }
 

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

@@ -246,7 +246,34 @@ suite('ExtensionService', () => {
 
  test('issue #152204: Remote extension host not disposed after closing vscode client', async () => {
  await extService.startExtensionHosts();
- extService.stopExtensionHosts();
+ extService.stopExtensionHosts(true);
  assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3', 'dispose 3', 'dispose 2', 'dispose 1']));
  });
+
+ test('Extension host disposed when awaited', async () => {
+ await extService.startExtensionHosts();
+ await extService.stopExtensionHosts();
+ assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3', 'dispose 3', 'dispose 2', 'dispose 1']));
+ });
+
+ test('Extension host not disposed when vetoed (sync)', async () => {
+ await extService.startExtensionHosts();
+
+ extService.onWillStop(e => e.veto(true, 'test 1'));
+ extService.onWillStop(e => e.veto(false, 'test 2'));
+
+ await extService.stopExtensionHosts();
+ assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3']));
+ });
+
+ test('Extension host not disposed when vetoed (async)', async () => {
+ await extService.startExtensionHosts();
+
+ extService.onWillStop(e => e.veto(false, 'test 1'));
+ extService.onWillStop(e => e.veto(Promise.resolve(true), 'test 2'));
+ extService.onWillStop(e => e.veto(Promise.resolve(false), 'test 3'));
+
+ await extService.stopExtensionHosts();
+ assert.deepStrictEqual(extService.order, (['create 1', 'create 2', 'create 3']));
+ });
 });

src/vs/workbench/services/userDataProfile/browser/userDataProfileManagement.ts

@@ -117,7 +117,7 @@ export class UserDataProfileManagementService extends Disposable implements IUse
  const isRemoteWindow = !!this.environmentService.remoteAuthority;
 
  if (!isRemoteWindow) {
- this.extensionService.stopExtensionHosts();
+ this.extensionService.stopExtensionHosts(true); // TODO@sandy081 adopt support for extension host to veto stopping
  }
 
  // In a remote window update current profile before reloading so that data is preserved from current profile if asked to preserve

src/vs/workbench/services/workspaces/electron-sandbox/workspaceEditingService.ts

@@ -154,6 +154,11 @@ export class NativeWorkspaceEditingService extends AbstractWorkspaceEditingServi
  }
 
  async enterWorkspace(workspaceUri: URI): Promise<void> {
+ const stopped = await this.extensionService.stopExtensionHosts();
+ if (!stopped) {
+ return;
+ }
+
  const result = await this.doEnterWorkspace(workspaceUri);
  if (result) {
 
@@ -175,7 +180,7 @@ export class NativeWorkspaceEditingService extends AbstractWorkspaceEditingServi
  // Restart the extension host: entering a workspace means a new location for
  // storage and potentially a change in the workspace.rootPath property.
  else {
- this.extensionService.restartExtensionHost();
+ this.extensionService.startExtensionHosts();
  }
  }
 }