Add DisposableMap helper (#163006)

Adds `DisposableMap` to help manage the lifecycle of maps of  disposable values. This is useful for a few reasons:

- `DisposableMap` is itself disposable, so you can use it with `_register` or add it to a `DisposableStore`
- The implementation of `set` on `DisposableMap` prevents leaking values on override
- The `delete` implementation also makes sure it disposes of the deleted values

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 {

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();
-	}
 }

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,7 +24,6 @@ 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';
@@ -32,15 +31,16 @@ import { CustomTextEditorModel } from 'vs/workbench/contrib/customEditor/common/
 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);
 	}

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 {

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();
 		}
 	}

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,9 +46,8 @@ 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> {
@@ -56,8 +55,7 @@ export class MainThreadFileSystem implements MainThreadFileSystemShape {
 	}
 
 	$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);
 		}
 	}
 }

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);
 	}
 }

src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts

@@ -8,7 +8,7 @@ import { CancellationToken } from 'vs/base/common/cancellation';
 import { createStringDataTransferItem, VSDataTransfer } from 'vs/base/common/dataTransfer';
 import { CancellationError } from 'vs/base/common/errors';
 import { Emitter, Event } from 'vs/base/common/event';
-import { combinedDisposable, Disposable, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
+import { combinedDisposable, Disposable, DisposableMap, toDisposable } from 'vs/base/common/lifecycle';
 import { revive } from 'vs/base/common/marshalling';
 import { mixin } from 'vs/base/common/objects';
 import { URI } from 'vs/base/common/uri';
@@ -38,7 +38,7 @@ import { ExtHostContext, ExtHostLanguageFeaturesShape, ICallHierarchyItemDto, IC
 export class MainThreadLanguageFeatures extends Disposable implements MainThreadLanguageFeaturesShape {
 
 	private readonly _proxy: ExtHostLanguageFeaturesShape;
-	private readonly _registrations = new Map<number, IDisposable>();
+	private readonly _registrations = this._register(new DisposableMap<number>());
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -80,21 +80,8 @@ export class MainThreadLanguageFeatures extends Disposable implements MainThread
 		}
 	}
 
-	override dispose(): void {
-		for (const registration of this._registrations.values()) {
-			registration.dispose();
-		}
-		this._registrations.clear();
-
-		super.dispose();
-	}
-
 	$unregister(handle: number): void {
-		const registration = this._registrations.get(handle);
-		if (registration) {
-			registration.dispose();
-			this._registrations.delete(handle);
-		}
+		this._registrations.deleteAndDispose(handle);
 	}
 
 	//#region --- revive functions

src/vs/workbench/api/browser/mainThreadLanguages.ts

@@ -13,7 +13,7 @@ import { IRange, Range } from 'vs/editor/common/core/range';
 import { StandardTokenType } from 'vs/editor/common/encodedTokenAttributes';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
 import { ILanguageStatus, ILanguageStatusService } from 'vs/workbench/services/languageStatus/common/languageStatusService';
-import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { DisposableMap, DisposableStore } from 'vs/base/common/lifecycle';
 
 @extHostNamedCustomer(MainContext.MainThreadLanguages)
 export class MainThreadLanguages implements MainThreadLanguagesShape {
@@ -21,7 +21,7 @@ export class MainThreadLanguages implements MainThreadLanguagesShape {
 	private readonly _disposables = new DisposableStore();
 	private readonly _proxy: ExtHostLanguagesShape;
 
-	private readonly _status = new Map<number, IDisposable>();
+	private readonly _status = new DisposableMap<number>();
 
 	constructor(
 		_extHostContext: IExtHostContext,
@@ -40,11 +40,7 @@ export class MainThreadLanguages implements MainThreadLanguagesShape {
 
 	dispose(): void {
 		this._disposables.dispose();
-
-		for (const status of this._status.values()) {
-			status.dispose();
-		}
-		this._status.clear();
+		this._status.dispose();
 	}
 
 	async $changeLanguage(resource: UriComponents, languageId: string): Promise<void> {

src/vs/workbench/api/browser/mainThreadNotebookDocumentsAndEditors.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { diffMaps, diffSets } from 'vs/base/common/collections';
-import { combinedDisposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { combinedDisposable, DisposableStore, DisposableMap } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { MainThreadNotebookDocuments } from 'vs/workbench/api/browser/mainThreadNotebookDocuments';
@@ -85,7 +85,7 @@ export class MainThreadNotebooksAndEditors {
 	private readonly _proxy: Pick<ExtHostNotebookShape, '$acceptDocumentAndEditorsDelta'>;
 	private readonly _disposables = new DisposableStore();
 
-	private readonly _editorListeners = new Map<string, IDisposable>();
+	private readonly _editorListeners = new DisposableMap<string>();
 
 	private _currentState?: NotebookAndEditorState;
 
@@ -121,6 +121,7 @@ export class MainThreadNotebooksAndEditors {
 		this._mainThreadNotebooks.dispose();
 		this._mainThreadEditors.dispose();
 		this._disposables.dispose();
+		this._editorListeners.dispose();
 	}
 
 	private _handleEditorAdd(editor: INotebookEditor): void {
@@ -132,8 +133,7 @@ export class MainThreadNotebooksAndEditors {
 	}
 
 	private _handleEditorRemove(editor: INotebookEditor): void {
-		this._editorListeners.get(editor.getId())?.dispose();
-		this._editorListeners.delete(editor.getId());
+		this._editorListeners.deleteAndDispose(editor.getId());
 		this._updateState();
 	}
 

src/vs/workbench/api/browser/mainThreadNotebookKernels.ts

@@ -6,7 +6,7 @@
 import { isNonEmptyArray } from 'vs/base/common/arrays';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { Emitter, Event } from 'vs/base/common/event';
-import { combinedDisposable, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
+import { combinedDisposable, DisposableMap, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { ILanguageService } from 'vs/editor/common/languages/language';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
@@ -104,7 +104,7 @@ abstract class MainThreadKernel implements INotebookKernel {
 @extHostNamedCustomer(MainContext.MainThreadNotebookKernels)
 export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape {
 
-	private readonly _editors = new Map<INotebookEditor, IDisposable>();
+	private readonly _editors = new DisposableMap<INotebookEditor>();
 	private readonly _disposables = new DisposableStore();
 
 	private readonly _kernels = new Map<number, [kernel: MainThreadKernel, registraion: IDisposable]>();
@@ -143,6 +143,7 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 		for (const [, registration] of this._kernels.values()) {
 			registration.dispose();
 		}
+		this._editors.dispose();
 	}
 
 	// --- kernel ipc
@@ -168,8 +169,7 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 	}
 
 	private _onEditorRemove(editor: INotebookEditor) {
-		this._editors.get(editor)?.dispose();
-		this._editors.delete(editor);
+		this._editors.deleteAndDispose(editor);
 	}
 
 	async $postMessage(handle: number, editorId: string | undefined, message: any): Promise<boolean> {

src/vs/workbench/api/browser/mainThreadWebviewPanels.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { Disposable, dispose, IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableMap } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
@@ -83,9 +83,7 @@ export class MainThreadWebviewPanels extends Disposable implements extHostProtoc
 
 	private readonly _webviewInputs = new WebviewInputStore();
 
-	private readonly _editorProviders = new Map<string, IDisposable>();
-
-	private readonly _revivers = new Map<string, IDisposable>();
+	private readonly _revivers = this._register(new DisposableMap<string>());
 
 	constructor(
 		context: IExtHostContext,
@@ -127,16 +125,6 @@ export class MainThreadWebviewPanels extends Disposable implements extHostProtoc
 		}));
 	}
 
-	override dispose() {
-		super.dispose();
-
-		dispose(this._editorProviders.values());
-		this._editorProviders.clear();
-
-		dispose(this._revivers.values());
-		this._revivers.clear();
-	}
-
 	public get webviewInputs(): Iterable<WebviewInput> { return this._webviewInputs; }
 
 	public addWebviewInput(handle: extHostProtocol.WebviewHandle, input: WebviewInput, options: { serializeBuffersForPostMessage: boolean }): void {
@@ -295,13 +283,11 @@ export class MainThreadWebviewPanels extends Disposable implements extHostProtoc
 	}
 
 	public $unregisterSerializer(viewType: string): void {
-		const reviver = this._revivers.get(viewType);
-		if (!reviver) {
+		if (!this._revivers.has(viewType)) {
 			throw new Error(`No reviver for ${viewType} registered`);
 		}
 
-		reviver.dispose();
-		this._revivers.delete(viewType);
+		this._revivers.deleteAndDispose(viewType);
 	}
 
 	private updateWebviewViewStates(activeEditorInput: EditorInput | undefined) {