microsoft · mjbvz · Oct 7, 2022 · Oct 7, 2022
diff --git a/src/vs/base/common/lifecycle.ts b/src/vs/base/common/lifecycle.ts
@@ -443,3 +443,51 @@ export function disposeOnReturn(fn: (store: DisposableStore) => void): void {
 		store.dispose();
 	}
 }
+
+/**
+ * A map the manages the lifecycle of the values that it stores.
+ */
+export class DisposableMap<K, V extends IDisposable = IDisposable> implements IDisposable {
+
+	private readonly _store = new Map<K, V>();
+	private _isDisposed = false;
+
+	dispose() {
+		this.clearAndDisposeAll();
+		this._isDisposed = true;
+	}
+
+	clearAndDisposeAll() {
+		dispose(this._store.values());
+		this._store.clear();
+	}
+
+	has(key: K): boolean {
+		return this._store.has(key);
+	}
+
+	get(key: K): V | undefined {
+		return this._store.get(key);
+	}
+
+	set(key: K, value: V, skipDisposeOnOverwrite = false): void {
+		if (this._isDisposed) {
+			console.warn(new Error('Trying to add a disposable to a DisposableMap that has already been disposed of. The added object will be leaked!').stack);
+		}
+
+		if (!skipDisposeOnOverwrite) {
+			this._store.get(key)?.dispose();
+		}
+
+		this._store.set(key, value);
+	}
+
+	deleteAndDispose(key: K): void {
+		this._store.get(key)?.dispose();
+		this._store.delete(key);
+	}
+
+	[Symbol.iterator](): IterableIterator<[K, V]> {
+		return this._store[Symbol.iterator]();
+	}
+}
diff --git a/src/vs/workbench/api/browser/mainThreadCommands.ts b/src/vs/workbench/api/browser/mainThreadCommands.ts
@@ -3,19 +3,19 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ICommandService, CommandsRegistry, ICommandHandlerDescription } from 'vs/platform/commands/common/commands';
-import { IDisposable, dispose } from 'vs/base/common/lifecycle';
-import { ExtHostContext, MainThreadCommandsShape, ExtHostCommandsShape, MainContext } from '../common/extHost.protocol';
-import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
+import { DisposableMap, IDisposable } from 'vs/base/common/lifecycle';
 import { revive } from 'vs/base/common/marshalling';
+import { CommandsRegistry, ICommandHandlerDescription, ICommandService } from 'vs/platform/commands/common/commands';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
-import { SerializableObjectWithBuffers, Dto } from 'vs/workbench/services/extensions/common/proxyIdentifier';
+import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
+import { Dto, SerializableObjectWithBuffers } from 'vs/workbench/services/extensions/common/proxyIdentifier';
+import { ExtHostCommandsShape, ExtHostContext, MainContext, MainThreadCommandsShape } from '../common/extHost.protocol';
 
 
 @extHostNamedCustomer(MainContext.MainThreadCommands)
 export class MainThreadCommands implements MainThreadCommandsShape {
 
-	private readonly _commandRegistrations = new Map<string, IDisposable>();
+	private readonly _commandRegistrations = new DisposableMap<string>();
 	private readonly _generateCommandsDocumentationRegistration: IDisposable;
 	private readonly _proxy: ExtHostCommandsShape;
 
@@ -30,9 +30,7 @@ export class MainThreadCommands implements MainThreadCommandsShape {
 	}
 
 	dispose() {
-		dispose(this._commandRegistrations.values());
-		this._commandRegistrations.clear();
-
+		this._commandRegistrations.dispose();
 		this._generateCommandsDocumentationRegistration.dispose();
 	}
 
@@ -67,11 +65,7 @@ export class MainThreadCommands implements MainThreadCommandsShape {
 	}
 
 	$unregisterCommand(id: string): void {
-		const command = this._commandRegistrations.get(id);
-		if (command) {
-			command.dispose();
-			this._commandRegistrations.delete(id);
-		}
+		this._commandRegistrations.deleteAndDispose(id);
 	}
 
 	$fireCommandActivationEvent(id: string): void {

diff --git a/src/vs/workbench/api/browser/mainThreadComments.ts b/src/vs/workbench/api/browser/mainThreadComments.ts
@@ -5,7 +5,7 @@
 
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { Emitter, Event } from 'vs/base/common/event';
-import { Disposable, DisposableStore, dispose, IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { generateUuid } from 'vs/base/common/uuid';
 import { IRange, Range } from 'vs/editor/common/core/range';
@@ -465,8 +465,7 @@ const commentsViewIcon = registerIcon('comments-view-icon', Codicon.commentDiscu
 @extHostNamedCustomer(MainContext.MainThreadComments)
 export class MainThreadComments extends Disposable implements MainThreadCommentsShape {
 	private readonly _proxy: ExtHostCommentsShape;
-	private _documentProviders = new Map<number, IDisposable>();
-	private _workspaceProviders = new Map<number, IDisposable>();
+
 	private _handlers = new Map<number, string>();
 	private _commentControllers = new Map<number, MainThreadCommentController>();
 
@@ -671,11 +670,4 @@ export class MainThreadComments extends Disposable implements MainThreadComments
 		}
 		return this._handlers.get(handle)!;
 	}
-	override dispose(): void {
-		super.dispose();
-		this._workspaceProviders.forEach(value => dispose(value));
-		this._workspaceProviders.clear();
-		this._documentProviders.forEach(value => dispose(value));
-		this._documentProviders.clear();
-	}
 }
diff --git a/src/vs/workbench/api/browser/mainThreadCustomEditors.ts b/src/vs/workbench/api/browser/mainThreadCustomEditors.ts
@@ -9,7 +9,7 @@ import { VSBuffer } from 'vs/base/common/buffer';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { isCancellationError, onUnexpectedError } from 'vs/base/common/errors';
 import { Emitter, Event } from 'vs/base/common/event';
-import { Disposable, DisposableStore, dispose, IDisposable, IReference } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableMap, DisposableStore, IReference } from 'vs/base/common/lifecycle';
 import { Schemas } from 'vs/base/common/network';
 import { basename } from 'vs/base/common/path';
 import { isEqual, isEqualOrParent, toLocalResource } from 'vs/base/common/resources';
@@ -24,23 +24,23 @@ import { MainThreadWebviewPanels } from 'vs/workbench/api/browser/mainThreadWebv
 import { MainThreadWebviews, reviveWebviewExtension } from 'vs/workbench/api/browser/mainThreadWebviews';
 import * as extHostProtocol from 'vs/workbench/api/common/extHost.protocol';
 import { IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
-import { editorGroupToColumn } from 'vs/workbench/services/editor/common/editorGroupColumn';
 import { CustomEditorInput } from 'vs/workbench/contrib/customEditor/browser/customEditorInput';
 import { CustomDocumentBackupData } from 'vs/workbench/contrib/customEditor/browser/customEditorInputFactory';
 import { ICustomEditorModel, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
 import { CustomTextEditorModel } from 'vs/workbench/contrib/customEditor/common/customTextEditorModel';
 import { WebviewExtensionDescription } from 'vs/workbench/contrib/webview/browser/webview';
 import { WebviewInput } from 'vs/workbench/contrib/webviewPanel/browser/webviewEditorInput';
 import { IWebviewWorkbenchService } from 'vs/workbench/contrib/webviewPanel/browser/webviewWorkbenchService';
+import { editorGroupToColumn } from 'vs/workbench/services/editor/common/editorGroupColumn';
 import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
+import { IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
 import { IPathService } from 'vs/workbench/services/path/common/pathService';
+import { ResourceWorkingCopy } from 'vs/workbench/services/workingCopy/common/resourceWorkingCopy';
+import { IWorkingCopy, IWorkingCopyBackup, IWorkingCopySaveEvent, NO_TYPE_ID, WorkingCopyCapabilities } from 'vs/workbench/services/workingCopy/common/workingCopy';
 import { IWorkingCopyFileService, WorkingCopyFileEvent } from 'vs/workbench/services/workingCopy/common/workingCopyFileService';
 import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
-import { IWorkingCopy, IWorkingCopyBackup, IWorkingCopySaveEvent, NO_TYPE_ID, WorkingCopyCapabilities } from 'vs/workbench/services/workingCopy/common/workingCopy';
-import { ResourceWorkingCopy } from 'vs/workbench/services/workingCopy/common/resourceWorkingCopy';
-import { IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
 
 const enum CustomEditorModelType {
 	Custom,
@@ -51,7 +51,7 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 
 	private readonly _proxyCustomEditors: extHostProtocol.ExtHostCustomEditorsShape;
 
-	private readonly _editorProviders = new Map<string, IDisposable>();
+	private readonly _editorProviders = this._register(new DisposableMap<string>());
 
 	private readonly _editorRenameBackups = new Map<string, CustomDocumentBackupData>();
 
@@ -99,13 +99,6 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 		this._register(workingCopyFileService.onWillRunWorkingCopyFileOperation(async e => this.onWillRunWorkingCopyFileOperation(e)));
 	}
 
-	override dispose() {
-		super.dispose();
-
-		dispose(this._editorProviders.values());
-		this._editorProviders.clear();
-	}
-
 	public $registerTextEditorProvider(extensionData: extHostProtocol.WebviewExtensionDescription, viewType: string, options: extHostProtocol.IWebviewPanelOptions, capabilities: extHostProtocol.CustomTextEditorCapabilities, serializeBuffersForPostMessage: boolean): void {
 		this.registerEditorProvider(CustomEditorModelType.Text, reviveWebviewExtension(extensionData), viewType, options, capabilities, true, serializeBuffersForPostMessage);
 	}
@@ -211,13 +204,11 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 	}
 
 	public $unregisterEditorProvider(viewType: string): void {
-		const provider = this._editorProviders.get(viewType);
-		if (!provider) {
+		if (!this._editorProviders.has(viewType)) {
 			throw new Error(`No provider for ${viewType} registered`);
 		}
 
-		provider.dispose();
-		this._editorProviders.delete(viewType);
+		this._editorProviders.deleteAndDispose(viewType);
 
 		this._customEditorService.models.disposeAllModelsForView(viewType);
 	}

diff --git a/src/vs/workbench/api/browser/mainThreadDocumentContentProviders.ts b/src/vs/workbench/api/browser/mainThreadDocumentContentProviders.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { IDisposable, dispose } from 'vs/base/common/lifecycle';
+import { dispose, DisposableMap } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { EditOperation } from 'vs/editor/common/core/editOperation';
 import { Range } from 'vs/editor/common/core/range';
@@ -20,7 +20,7 @@ import { CancellationTokenSource } from 'vs/base/common/cancellation';
 @extHostNamedCustomer(MainContext.MainThreadDocumentContentProviders)
 export class MainThreadDocumentContentProviders implements MainThreadDocumentContentProvidersShape {
 
-	private readonly _resourceContentProvider = new Map<number, IDisposable>();
+	private readonly _resourceContentProvider = new DisposableMap<number>();
 	private readonly _pendingUpdate = new Map<string, CancellationTokenSource>();
 	private readonly _proxy: ExtHostDocumentContentProvidersShape;
 
@@ -35,7 +35,7 @@ export class MainThreadDocumentContentProviders implements MainThreadDocumentCon
 	}
 
 	dispose(): void {
-		dispose(this._resourceContentProvider.values());
+		this._resourceContentProvider.dispose();
 		dispose(this._pendingUpdate.values());
 	}
 
@@ -56,11 +56,7 @@ export class MainThreadDocumentContentProviders implements MainThreadDocumentCon
 	}
 
 	$unregisterTextContentProvider(handle: number): void {
-		const registration = this._resourceContentProvider.get(handle);
-		if (registration) {
-			registration.dispose();
-			this._resourceContentProvider.delete(handle);
-		}
+		this._resourceContentProvider.deleteAndDispose(handle);
 	}
 
 	$onVirtualDocumentChange(uri: UriComponents, value: string): void {

diff --git a/src/vs/workbench/api/browser/mainThreadDocumentsAndEditors.ts b/src/vs/workbench/api/browser/mainThreadDocumentsAndEditors.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Event } from 'vs/base/common/event';
-import { IDisposable, combinedDisposable, DisposableStore } from 'vs/base/common/lifecycle';
+import { combinedDisposable, DisposableStore, DisposableMap } from 'vs/base/common/lifecycle';
 import { ICodeEditor, isCodeEditor, isDiffEditor, IActiveCodeEditor } from 'vs/editor/browser/editorBrowser';
 import { ICodeEditorService } from 'vs/editor/browser/services/codeEditorService';
 import { IEditor } from 'vs/editor/common/editorCommon';
@@ -113,7 +113,7 @@ const enum ActiveEditorOrder {
 class MainThreadDocumentAndEditorStateComputer {
 
 	private readonly _toDispose = new DisposableStore();
-	private _toDisposeOnEditorRemove = new Map<string, IDisposable>();
+	private readonly _toDisposeOnEditorRemove = new DisposableMap<string>();
 	private _currentState?: DocumentAndEditorState;
 	private _activeEditorOrder: ActiveEditorOrder = ActiveEditorOrder.Editor;
 
@@ -141,6 +141,7 @@ class MainThreadDocumentAndEditorStateComputer {
 
 	dispose(): void {
 		this._toDispose.dispose();
+		this._toDisposeOnEditorRemove.dispose();
 	}
 
 	private _onDidAddEditor(e: ICodeEditor): void {
@@ -153,10 +154,9 @@ class MainThreadDocumentAndEditorStateComputer {
 	}
 
 	private _onDidRemoveEditor(e: ICodeEditor): void {
-		const sub = this._toDisposeOnEditorRemove.get(e.getId());
-		if (sub) {
-			this._toDisposeOnEditorRemove.delete(e.getId());
-			sub.dispose();
+		const id = e.getId();
+		if (this._toDisposeOnEditorRemove.has(id)) {
+			this._toDisposeOnEditorRemove.deleteAndDispose(id);
 			this._updateState();
 		}
 	}

diff --git a/src/vs/workbench/api/browser/mainThreadFileSystem.ts b/src/vs/workbench/api/browser/mainThreadFileSystem.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Emitter, Event } from 'vs/base/common/event';
-import { IDisposable, dispose, toDisposable, DisposableStore } from 'vs/base/common/lifecycle';
+import { IDisposable, toDisposable, DisposableStore, DisposableMap } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { IFileWriteOptions, FileSystemProviderCapabilities, IFileChange, IFileService, IStat, IWatchOptions, FileType, IFileOverwriteOptions, IFileDeleteOptions, IFileOpenOptions, FileOperationError, FileOperationResult, FileSystemProviderErrorCode, IFileSystemProviderWithOpenReadWriteCloseCapability, IFileSystemProviderWithFileReadWriteCapability, IFileSystemProviderWithFileFolderCopyCapability, FilePermission, toFileSystemProviderErrorCode, IFilesConfiguration, IFileStatWithPartialMetadata, IFileStat } from 'vs/platform/files/common/files';
 import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
@@ -22,9 +22,9 @@ import { rtrim } from 'vs/base/common/strings';
 export class MainThreadFileSystem implements MainThreadFileSystemShape {
 
 	private readonly _proxy: ExtHostFileSystemShape;
-	private readonly _fileProvider = new Map<number, RemoteFileSystemProvider>();
+	private readonly _fileProvider = new DisposableMap<number, RemoteFileSystemProvider>();
 	private readonly _disposables = new DisposableStore();
-	private readonly _watches = new Map<number, IDisposable>();
+	private readonly _watches = new DisposableMap<number>();
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -46,18 +46,16 @@ export class MainThreadFileSystem implements MainThreadFileSystemShape {
 
 	dispose(): void {
 		this._disposables.dispose();
-		dispose(this._fileProvider.values());
-		dispose(this._watches.values());
-		this._fileProvider.clear();
+		this._fileProvider.dispose();
+		this._watches.dispose();
 	}
 
 	async $registerFileSystemProvider(handle: number, scheme: string, capabilities: FileSystemProviderCapabilities): Promise<void> {
 		this._fileProvider.set(handle, new RemoteFileSystemProvider(this._fileService, scheme, capabilities, handle, this._proxy));
 	}
 
 	$unregisterProvider(handle: number): void {
-		this._fileProvider.get(handle)?.dispose();
-		this._fileProvider.delete(handle);
+		this._fileProvider.deleteAndDispose(handle);
 	}
 
 	$onFileSystemChange(handle: number, changes: IFileChangeDto[]): void {
@@ -254,12 +252,9 @@ export class MainThreadFileSystem implements MainThreadFileSystemShape {
 	}
 
 	$unwatch(session: number): void {
-		const subscription = this._watches.get(session);
-		if (subscription) {
+		if (this._watches.has(session)) {
 			this._logService.trace(`MainThreadFileSystem#$unwatch(): request to stop watching (session: ${session})`);
-
-			subscription.dispose();
-			this._watches.delete(session);
+			this._watches.deleteAndDispose(session);
 		}
 	}
 }

diff --git a/src/vs/workbench/api/browser/mainThreadLabelService.ts b/src/vs/workbench/api/browser/mainThreadLabelService.ts
@@ -3,20 +3,22 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { Disposable, DisposableMap } from 'vs/base/common/lifecycle';
+import { ILabelService, ResourceLabelFormatter } from 'vs/platform/label/common/label';
 import { MainContext, MainThreadLabelServiceShape } from 'vs/workbench/api/common/extHost.protocol';
 import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
-import { ResourceLabelFormatter, ILabelService } from 'vs/platform/label/common/label';
-import { IDisposable, dispose } from 'vs/base/common/lifecycle';
 
 @extHostNamedCustomer(MainContext.MainThreadLabelService)
-export class MainThreadLabelService implements MainThreadLabelServiceShape {
+export class MainThreadLabelService extends Disposable implements MainThreadLabelServiceShape {
 
-	private readonly _resourceLabelFormatters = new Map<number, IDisposable>();
+	private readonly _resourceLabelFormatters = this._register(new DisposableMap<number>());
 
 	constructor(
 		_: IExtHostContext,
 		@ILabelService private readonly _labelService: ILabelService
-	) { }
+	) {
+		super();
+	}
 
 	$registerResourceLabelFormatter(handle: number, formatter: ResourceLabelFormatter): void {
 		// Dynamicily registered formatters should have priority over those contributed via package.json
@@ -26,11 +28,6 @@ export class MainThreadLabelService implements MainThreadLabelServiceShape {
 	}
 
 	$unregisterResourceLabelFormatter(handle: number): void {
-		dispose(this._resourceLabelFormatters.get(handle));
-		this._resourceLabelFormatters.delete(handle);
-	}
-
-	dispose(): void {
-		// noop
+		this._resourceLabelFormatters.deleteAndDispose(handle);
 	}
 }