Debt - do not expose an unregister method from service (#158641)

* Debt - do not expose an unregister method from service * Add registration disposable to disposable store
microsoft · Aug 30, 2022 · db4f5a8 · db4f5a8
1 parent af14c10
commit db4f5a8
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 16 deletions.
diff --git a/src/vs/platform/workspace/common/editSessions.ts b/src/vs/platform/workspace/common/editSessions.ts
@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
+import { IDisposable } from 'vs/base/common/lifecycle';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { IWorkspaceFolder } from 'vs/platform/workspace/common/workspace';
 
@@ -17,7 +18,6 @@ export const IEditSessionIdentityService = createDecorator<IEditSessionIdentityS
 export interface IEditSessionIdentityService {
 	readonly _serviceBrand: undefined;
 
-	registerEditSessionIdentityProvider(provider: IEditSessionIdentityProvider): void;
-	unregisterEditSessionIdentityProvider(scheme: string): void;
+	registerEditSessionIdentityProvider(provider: IEditSessionIdentityProvider): IDisposable;
 	getEditSessionIdentifier(workspaceFolder: IWorkspaceFolder, cancellationTokenSource: CancellationTokenSource): Promise<string | undefined>;
 }
diff --git a/src/vs/workbench/api/browser/mainThreadWorkspace.ts b/src/vs/workbench/api/browser/mainThreadWorkspace.ts
@@ -5,7 +5,7 @@
 
 import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { isCancellationError } from 'vs/base/common/errors';
-import { DisposableStore } from 'vs/base/common/lifecycle';
+import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { isNative } from 'vs/base/common/platform';
 import { withNullAsUndefined } from 'vs/base/common/types';
 import { URI, UriComponents } from 'vs/base/common/uri';
@@ -224,16 +224,23 @@ export class MainThreadWorkspace implements MainThreadWorkspaceShape {
 	}
 
 	// --- edit sessions ---
-	$registerEditSessionIdentityProvider(scheme: string) {
-		this._editSessionIdentityService.registerEditSessionIdentityProvider({
+	private registeredEditSessionProviders = new Map<number, IDisposable>();
+
+	$registerEditSessionIdentityProvider(handle: number, scheme: string) {
+		const disposable = this._editSessionIdentityService.registerEditSessionIdentityProvider({
 			scheme: scheme,
 			getEditSessionIdentifier: async (workspaceFolder: WorkspaceFolder, token: CancellationToken) => {
 				return this._proxy.$getEditSessionIdentifier(workspaceFolder.uri, token);
 			}
 		});
+
+		this.registeredEditSessionProviders.set(handle, disposable);
+		this._toDispose.add(disposable);
 	}
 
-	$unregisterEditSessionIdentityProvider(scheme: string) {
-		this._editSessionIdentityService.unregisterEditSessionIdentityProvider(scheme);
+	$unregisterEditSessionIdentityProvider(handle: number) {
+		const disposable = this.registeredEditSessionProviders.get(handle);
+		disposable?.dispose();
+		this.registeredEditSessionProviders.delete(handle);
 	}
 }
diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts
@@ -1085,8 +1085,8 @@ export interface MainThreadWorkspaceShape extends IDisposable {
 	$updateWorkspaceFolders(extensionName: string, index: number, deleteCount: number, workspaceFoldersToAdd: { uri: UriComponents; name?: string }[]): Promise<void>;
 	$resolveProxy(url: string): Promise<string | undefined>;
 	$requestWorkspaceTrust(options?: WorkspaceTrustRequestOptions): Promise<boolean | undefined>;
-	$registerEditSessionIdentityProvider(scheme: string): void;
-	$unregisterEditSessionIdentityProvider(scheme: string): void;
+	$registerEditSessionIdentityProvider(handle: number, scheme: string): void;
+	$unregisterEditSessionIdentityProvider(handle: number): void;
 }
 
 export interface IFileChangeDto {

diff --git a/src/vs/workbench/api/common/extHostWorkspace.ts b/src/vs/workbench/api/common/extHostWorkspace.ts
@@ -583,6 +583,8 @@ export class ExtHostWorkspace implements ExtHostWorkspaceShape, IExtHostWorkspac
 
 	// --- edit sessions ---
 
+	private _providerHandlePool = 0;
+
 	// called by ext host
 	registerEditSessionIdentityProvider(scheme: string, provider: vscode.EditSessionIdentityProvider) {
 		if (this._editSessionIdentityProviders.has(scheme)) {
@@ -591,11 +593,12 @@ export class ExtHostWorkspace implements ExtHostWorkspaceShape, IExtHostWorkspac
 
 		this._editSessionIdentityProviders.set(scheme, provider);
 		const outgoingScheme = this._uriTransformerService.transformOutgoingScheme(scheme);
-		this._proxy.$registerEditSessionIdentityProvider(outgoingScheme);
+		const handle = this._providerHandlePool++;
+		this._proxy.$registerEditSessionIdentityProvider(handle, outgoingScheme);
 
 		return toDisposable(() => {
 			this._editSessionIdentityProviders.delete(scheme);
-			this._proxy.$unregisterEditSessionIdentityProvider(outgoingScheme);
+			this._proxy.$unregisterEditSessionIdentityProvider(handle);
 		});
 	}
 

diff --git a/src/vs/workbench/services/workspaces/common/editSessionIdentityService.ts b/src/vs/workbench/services/workspaces/common/editSessionIdentityService.ts
@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
+import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { ILogService } from 'vs/platform/log/common/log';
 import { IEditSessionIdentityProvider, IEditSessionIdentityService } from 'vs/platform/workspace/common/editSessions';
@@ -20,16 +21,15 @@ export class EditSessionIdentityService implements IEditSessionIdentityService {
 		@ILogService private readonly _logService: ILogService,
 	) { }
 
-	registerEditSessionIdentityProvider(provider: IEditSessionIdentityProvider): void {
+	registerEditSessionIdentityProvider(provider: IEditSessionIdentityProvider): IDisposable {
 		if (this._editSessionIdentifierProviders.get(provider.scheme)) {
 			throw new Error(`A provider has already been registered for scheme ${provider.scheme}`);
 		}
 
 		this._editSessionIdentifierProviders.set(provider.scheme, provider);
-	}
-
-	unregisterEditSessionIdentityProvider(scheme: string): void {
-		this._editSessionIdentifierProviders.delete(scheme);
+		return toDisposable(() => {
+			this._editSessionIdentifierProviders.delete(provider.scheme);
+		});
 	}
 
 	async getEditSessionIdentifier(workspaceFolder: IWorkspaceFolder, cancellationTokenSource: CancellationTokenSource): Promise<string | undefined> {