re #212879. dispose child instantiation service properly (#232532)

* re #212879. dispose child instantiation service properly * fix unit tests * Fix memory leaks
microsoft · Oct 29, 2024 · db1b27a · db1b27a
1 parent fcae80e
commit db1b27a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/src/vs/workbench/contrib/notebook/browser/services/notebookEditorServiceImpl.ts b/src/vs/workbench/contrib/notebook/browser/services/notebookEditorServiceImpl.ts
@@ -38,7 +38,7 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 	readonly onDidAddNotebookEditor = this._onNotebookEditorAdd.event;
 	readonly onDidRemoveNotebookEditor = this._onNotebookEditorsRemove.event;
 
-	private readonly _borrowableEditors = new Map<number, ResourceMap<{ widget: NotebookEditorWidget; editorType: string; token: number | undefined }[]>>();
+	private readonly _borrowableEditors = new Map<number, ResourceMap<{ widget: NotebookEditorWidget; editorType: string; token: number | undefined; disposableStore: DisposableStore }[]>>();
 
 	constructor(
 		@IEditorGroupsService private readonly editorGroupService: IEditorGroupsService,
@@ -65,6 +65,7 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 					const value = widgets.splice(index, 1)[0];
 					value.token = undefined;
 					this._disposeWidget(value.widget);
+					value.disposableStore.dispose();
 					value.widget = (<any>undefined); // unset the widget so that others that still hold a reference don't harm us
 				});
 			}));
@@ -98,6 +99,7 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 					for (const value of values) {
 						value.token = undefined;
 						this._disposeWidget(value.widget);
+						value.disposableStore.dispose();
 					}
 				}
 			}
@@ -125,6 +127,11 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 			listeners.forEach(listener => listener.dispose());
 		});
 		this.groupListener.clear();
+		this._borrowableEditors.forEach(widgetMap => {
+			widgetMap.forEach(widgets => {
+				widgets.forEach(widget => widget.disposableStore.dispose());
+			});
+		});
 	}
 
 	// --- group-based editor borrowing...
@@ -205,9 +212,10 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 			// NEW widget
 			const editorGroupContextKeyService = accessor.get(IContextKeyService);
 			const editorGroupEditorProgressService = accessor.get(IEditorProgressService);
-			const widget = this.createWidget(editorGroupContextKeyService, editorGroupEditorProgressService, creationOptions, codeWindow, initialDimension);
+			const widgetDisposeStore = new DisposableStore();
+			const widget = this.createWidget(editorGroupContextKeyService, widgetDisposeStore, editorGroupEditorProgressService, creationOptions, codeWindow, initialDimension);
 			const token = this._tokenPool++;
-			value = { widget, editorType: input.typeId, token };
+			value = { widget, editorType: input.typeId, token, disposableStore: widgetDisposeStore };
 
 			let map = this._borrowableEditors.get(groupId);
 			if (!map) {
@@ -217,7 +225,6 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 			const values = map.get(input.resource) ?? [];
 			values.push(value);
 			map.set(input.resource, values);
-
 		} else {
 			// reuse a widget which was either free'ed before or which
 			// is simply being reused...
@@ -228,10 +235,10 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 	}
 
 	// protected for unit testing overrides
-	protected createWidget(editorGroupContextKeyService: IContextKeyService, editorGroupEditorProgressService: IEditorProgressService, creationOptions?: INotebookEditorCreationOptions, codeWindow?: CodeWindow, initialDimension?: Dimension) {
-		const notebookInstantiationService = this.instantiationService.createChild(new ServiceCollection(
+	protected createWidget(editorGroupContextKeyService: IContextKeyService, widgetDisposeStore: DisposableStore, editorGroupEditorProgressService: IEditorProgressService, creationOptions?: INotebookEditorCreationOptions, codeWindow?: CodeWindow, initialDimension?: Dimension) {
+		const notebookInstantiationService = widgetDisposeStore.add(this.instantiationService.createChild(new ServiceCollection(
 			[IContextKeyService, editorGroupContextKeyService],
-			[IEditorProgressService, editorGroupEditorProgressService]));
+			[IEditorProgressService, editorGroupEditorProgressService])));
 		const ctorOptions = creationOptions ?? getDefaultNotebookCreationOptions();
 		const widget = notebookInstantiationService.createInstance(NotebookEditorWidget, {
 			...ctorOptions,

diff --git a/src/vs/workbench/contrib/notebook/test/browser/NotebookEditorWidgetService.test.ts b/src/vs/workbench/contrib/notebook/test/browser/NotebookEditorWidgetService.test.ts
@@ -135,6 +135,8 @@ suite('NotebookEditorWidgetService', () => {
 		assert.strictEqual(widget.value, undefined, 'widgets in group should get disposed');
 		assert.strictEqual(widgetDiffType.value, undefined, 'widgets in group should get disposed');
 		assert.notStrictEqual(widgetDiffGroup.value, undefined, 'other group should not be disposed');
+
+		notebookEditorService.dispose();
 	});
 
 	test('Widget should move between groups when editor is moved', async function () {