Disable notebook backups when file size is too large (microsoft#210493)

* configurable limit for notebook output size for backups

* disable backups entirely for notebooks with large outputs

* working copy - make the snapshot context more explicit

* test that snapshot will throw

* fix casing

---------

Co-authored-by: Benjamin Pasero <[email protected]>

Author: amunger

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

@@ -1096,5 +1096,10 @@ configurationRegistry.registerConfiguration({
 			type: 'boolean',
 			default: true
 		},
+		[NotebookSetting.outputBackupSizeLimit]: {
+			markdownDescription: nls.localize('notebook.backup.sizeLimit', "The limit of notebook output size in kilobytes (KB) where notebook files will no longer be backed up for hot reload. Use 0 for unlimited."),
+			type: 'number',
+			default: 10000
+		}
 	}
 });

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -962,6 +962,7 @@ export const NotebookSetting = {
 	notebookVariablesView: 'notebook.experimental.variablesView',
 	InteractiveWindowPromptToSave: 'interactiveWindow.promptToSaveOnClose',
 	cellFailureDiagnostics: 'notebook.cellFailureDiagnostics',
+	outputBackupSizeLimit: 'notebook.backup.sizeLimit',
 } as const;
 
 export const enum CellStatusbarAlignment {

src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts

@@ -21,7 +21,7 @@ import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/no
 import { ICellDto2, INotebookEditorModel, INotebookLoadOptions, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookData, NotebookSetting } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookSerializer, INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { IFilesConfigurationService } from 'vs/workbench/services/filesConfiguration/common/filesConfigurationService';
-import { IFileWorkingCopyModelConfiguration } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
+import { IFileWorkingCopyModelConfiguration, SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 import { IFileWorkingCopyManager } from 'vs/workbench/services/workingCopy/common/fileWorkingCopyManager';
 import { IStoredFileWorkingCopy, IStoredFileWorkingCopyModel, IStoredFileWorkingCopyModelContentChangedEvent, IStoredFileWorkingCopyModelFactory, IStoredFileWorkingCopySaveEvent, StoredFileWorkingCopyState } from 'vs/workbench/services/workingCopy/common/storedFileWorkingCopy';
 import { IUntitledFileWorkingCopy, IUntitledFileWorkingCopyModel, IUntitledFileWorkingCopyModelContentChangedEvent, IUntitledFileWorkingCopyModelFactory } from 'vs/workbench/services/workingCopy/common/untitledFileWorkingCopy';
@@ -252,14 +252,15 @@ export class NotebookFileWorkingCopyModel extends Disposable implements IStoredF
 		return this._notebookModel;
 	}
 
-	async snapshot(token: CancellationToken): Promise<VSBufferReadableStream> {
+	async snapshot(context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream> {
 		const serializer = await this.getNotebookSerializer();
 
 		const data: NotebookData = {
 			metadata: filter(this._notebookModel.metadata, key => !serializer.options.transientDocumentMetadata[key]),
 			cells: [],
 		};
 
+		let outputSize = 0;
 		for (const cell of this._notebookModel.cells) {
 			const cellData: ICellDto2 = {
 				cellKind: cell.cellKind,
@@ -270,6 +271,18 @@ export class NotebookFileWorkingCopyModel extends Disposable implements IStoredF
 				internalMetadata: cell.internalMetadata
 			};
 
+			const outputSizeLimit = this._configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;
+			if (context === SnapshotContext.Backup && outputSizeLimit > 0) {
+				cell.outputs.forEach(output => {
+					output.outputs.forEach(item => {
+						outputSize += item.data.byteLength;
+					});
+				});
+				if (outputSize > outputSizeLimit) {
+					throw new Error('Notebook too large to backup');
+				}
+			}
+
 			cellData.outputs = !serializer.options.transientOutputs ? cell.outputs : [];
 			cellData.metadata = filter(cell.metadata, key => !serializer.options.transientCellMetadata[key]);
 

src/vs/workbench/contrib/notebook/test/browser/notebookEditorModel.test.ts

@@ -15,10 +15,11 @@ import { TestConfigurationService } from 'vs/platform/configuration/test/common/
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
-import { CellKind, NotebookData, TransientOptions } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellKind, IOutputDto, NotebookData, NotebookSetting, TransientOptions } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookFileWorkingCopyModel } from 'vs/workbench/contrib/notebook/common/notebookEditorModel';
 import { INotebookSerializer, INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { setupInstantiationService } from 'vs/workbench/contrib/notebook/test/browser/testNotebookEditor';
+import { SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 
 suite('NotebookFileWorkingCopyModel', function () {
 
@@ -63,7 +64,7 @@ suite('NotebookFileWorkingCopyModel', function () {
 				configurationService
 			));
 
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 
@@ -84,7 +85,7 @@ suite('NotebookFileWorkingCopyModel', function () {
 				),
 				configurationService
 			));
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 	});
@@ -119,7 +120,7 @@ suite('NotebookFileWorkingCopyModel', function () {
 				configurationService
 			));
 
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 
@@ -140,7 +141,7 @@ suite('NotebookFileWorkingCopyModel', function () {
 				),
 				configurationService
 			));
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 	});
@@ -174,7 +175,7 @@ suite('NotebookFileWorkingCopyModel', function () {
 				configurationService
 			));
 
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 
@@ -195,10 +196,52 @@ suite('NotebookFileWorkingCopyModel', function () {
 				),
 				configurationService
 			));
-			await model.snapshot(CancellationToken.None);
+			await model.snapshot(SnapshotContext.Save, CancellationToken.None);
 			assert.strictEqual(callCount, 1);
 		}
 	});
+
+	test('Notebooks with outputs beyond the size threshold will throw for backup snapshots', async function () {
+		const outputLimit = 100;
+		await configurationService.setUserConfiguration(NotebookSetting.outputBackupSizeLimit, outputLimit * 1.0 / 1024);
+		const largeOutput: IOutputDto = { outputId: '123', outputs: [{ mime: Mimes.text, data: VSBuffer.fromString('a'.repeat(outputLimit + 1)) }] };
+		const notebook = instantiationService.createInstance(NotebookTextModel,
+			'notebook',
+			URI.file('test'),
+			[{ cellKind: CellKind.Code, language: 'foo', mime: 'foo', source: 'foo', outputs: [largeOutput], metadata: { foo: 123, bar: 456 } }],
+			{},
+			{ transientCellMetadata: {}, transientDocumentMetadata: {}, cellContentMetadata: {}, transientOutputs: false, }
+		);
+		disposables.add(notebook);
+
+		let callCount = 0;
+		const model = disposables.add(new NotebookFileWorkingCopyModel(
+			notebook,
+			mockNotebookService(notebook,
+				new class extends mock<INotebookSerializer>() {
+					override options: TransientOptions = { transientOutputs: true, transientDocumentMetadata: {}, transientCellMetadata: { bar: true }, cellContentMetadata: {} };
+					override async notebookToData(notebook: NotebookData) {
+						callCount += 1;
+						assert.strictEqual(notebook.cells[0].metadata!.foo, 123);
+						assert.strictEqual(notebook.cells[0].metadata!.bar, undefined);
+						return VSBuffer.fromString('');
+					}
+				}
+			),
+			configurationService
+		));
+
+		try {
+			await model.snapshot(SnapshotContext.Backup, CancellationToken.None);
+			assert.fail('Expected snapshot to throw an error for large output');
+		} catch (e) {
+			assert.notEqual(e.code, 'ERR_ASSERTION', e.message);
+		}
+
+		await model.snapshot(SnapshotContext.Save, CancellationToken.None);
+		assert.strictEqual(callCount, 1);
+
+	});
 });
 
 function mockNotebookService(notebook: NotebookTextModel, notebookSerializer: INotebookSerializer) {

src/vs/workbench/services/workingCopy/common/fileWorkingCopy.ts

@@ -35,6 +35,11 @@ export interface IFileWorkingCopyModelConfiguration {
 	readonly backupDelay?: number;
 }
 
+export const enum SnapshotContext {
+	Save = 1,
+	Backup = 2
+}
+
 /**
  * A generic file working copy model to be reused by untitled
  * and stored file working copies.
@@ -72,9 +77,10 @@ export interface IFileWorkingCopyModel extends IDisposable {
 	 * Snapshots the model's current content for writing. This must include
 	 * any changes that were made to the model that are in memory.
 	 *
+	 * @param context indicates in what context the snapshot is used
 	 * @param token support for cancellation
 	 */
-	snapshot(token: CancellationToken): Promise<VSBufferReadableStream>;
+	snapshot(context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream>;
 
 	/**
 	 * Updates the model with the provided contents. The implementation should

src/vs/workbench/services/workingCopy/common/fileWorkingCopyManager.ts

@@ -23,7 +23,7 @@ import { IUntitledFileWorkingCopy, IUntitledFileWorkingCopyModel, IUntitledFileW
 import { INewOrExistingUntitledFileWorkingCopyOptions, INewUntitledFileWorkingCopyOptions, INewUntitledFileWorkingCopyWithAssociatedResourceOptions, IUntitledFileWorkingCopyManager, UntitledFileWorkingCopyManager } from 'vs/workbench/services/workingCopy/common/untitledFileWorkingCopyManager';
 import { IWorkingCopyFileService } from 'vs/workbench/services/workingCopy/common/workingCopyFileService';
 import { IBaseFileWorkingCopyManager } from 'vs/workbench/services/workingCopy/common/abstractFileWorkingCopyManager';
-import { IFileWorkingCopy } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
+import { IFileWorkingCopy, SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 import { ILabelService } from 'vs/platform/label/common/label';
 import { ILogService } from 'vs/platform/log/common/log';
 import { INotificationService } from 'vs/platform/notification/common/notification';
@@ -390,7 +390,7 @@ export class FileWorkingCopyManager<S extends IStoredFileWorkingCopyModel, U ext
 		// use that to copy the contents to the target destination
 		const sourceWorkingCopy = this.get(source);
 		if (sourceWorkingCopy?.isResolved()) {
-			sourceContents = await sourceWorkingCopy.model.snapshot(CancellationToken.None);
+			sourceContents = await sourceWorkingCopy.model.snapshot(SnapshotContext.Save, CancellationToken.None);
 		}
 
 		// Otherwise we resolve the contents from the underlying file

src/vs/workbench/services/workingCopy/common/storedFileWorkingCopy.ts

@@ -27,7 +27,7 @@ import { IWorkingCopyEditorService } from 'vs/workbench/services/workingCopy/com
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IElevatedFileService } from 'vs/workbench/services/files/common/elevatedFileService';
 import { IResourceWorkingCopy, ResourceWorkingCopy } from 'vs/workbench/services/workingCopy/common/resourceWorkingCopy';
-import { IFileWorkingCopy, IFileWorkingCopyModel, IFileWorkingCopyModelFactory } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
+import { IFileWorkingCopy, IFileWorkingCopyModel, IFileWorkingCopyModelFactory, SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 import { IMarkdownString } from 'vs/base/common/htmlContent';
 
 /**
@@ -813,7 +813,7 @@ export class StoredFileWorkingCopy<M extends IStoredFileWorkingCopyModel> extend
 		// Fill in content if we are resolved
 		let content: VSBufferReadableStream | undefined = undefined;
 		if (this.isResolved()) {
-			content = await raceCancellation(this.model.snapshot(token), token);
+			content = await raceCancellation(this.model.snapshot(SnapshotContext.Backup, token), token);
 		}
 
 		return { meta, content };
@@ -1026,7 +1026,7 @@ export class StoredFileWorkingCopy<M extends IStoredFileWorkingCopyModel> extend
 					else {
 
 						// Snapshot working copy model contents
-						const snapshot = await raceCancellation(resolvedFileWorkingCopy.model.snapshot(saveCancellation.token), saveCancellation.token);
+						const snapshot = await raceCancellation(resolvedFileWorkingCopy.model.snapshot(SnapshotContext.Save, saveCancellation.token), saveCancellation.token);
 
 						// It is possible that a subsequent save is cancelling this
 						// running save. As such we return early when we detect that

src/vs/workbench/services/workingCopy/common/storedFileWorkingCopyManager.ts

@@ -29,6 +29,7 @@ import { IWorkingCopyEditorService } from 'vs/workbench/services/workingCopy/com
 import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
 import { isWeb } from 'vs/base/common/platform';
 import { onUnexpectedError } from 'vs/base/common/errors';
+import { SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 
 /**
  * The only one that should be dealing with `IStoredFileWorkingCopy` and handle all
@@ -352,7 +353,7 @@ export class StoredFileWorkingCopyManager<M extends IStoredFileWorkingCopyModel>
 							workingCopiesToRestore.push({
 								source: sourceResource,
 								target: targetResource,
-								snapshot: sourceWorkingCopy.isDirty() ? await sourceWorkingCopy.model?.snapshot(CancellationToken.None) : undefined
+								snapshot: sourceWorkingCopy.isDirty() ? await sourceWorkingCopy.model?.snapshot(SnapshotContext.Save, CancellationToken.None) : undefined
 							});
 						}
 					}

src/vs/workbench/services/workingCopy/common/untitledFileWorkingCopy.ts

@@ -6,7 +6,7 @@
 import { Event, Emitter } from 'vs/base/common/event';
 import { VSBufferReadableStream } from 'vs/base/common/buffer';
 import { IWorkingCopyBackup, IWorkingCopySaveEvent, WorkingCopyCapabilities } from 'vs/workbench/services/workingCopy/common/workingCopy';
-import { IFileWorkingCopy, IFileWorkingCopyModel, IFileWorkingCopyModelFactory } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
+import { IFileWorkingCopy, IFileWorkingCopyModel, IFileWorkingCopyModelFactory, SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
@@ -264,7 +264,7 @@ export class UntitledFileWorkingCopy<M extends IUntitledFileWorkingCopyModel> ex
 		// if any - to prevent backing up an unresolved working
 		// copy and loosing the initial value.
 		if (this.isResolved()) {
-			content = await raceCancellation(this.model.snapshot(token), token);
+			content = await raceCancellation(this.model.snapshot(SnapshotContext.Backup, token), token);
 		} else if (this.initialContents) {
 			content = this.initialContents.value;
 		}

src/vs/workbench/services/workingCopy/test/browser/storedFileWorkingCopy.test.ts

@@ -19,6 +19,7 @@ import { Promises, timeout } from 'vs/base/common/async';
 import { consumeReadable, consumeStream, isReadableStream } from 'vs/base/common/stream';
 import { runWithFakedTimers } from 'vs/base/test/common/timeTravelScheduler';
 import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
+import { SnapshotContext } from 'vs/workbench/services/workingCopy/common/fileWorkingCopy';
 
 export class TestStoredFileWorkingCopyModel extends Disposable implements IStoredFileWorkingCopyModel {
 
@@ -45,7 +46,7 @@ export class TestStoredFileWorkingCopyModel extends Disposable implements IStore
 		this.throwOnSnapshot = true;
 	}
 
-	async snapshot(token: CancellationToken): Promise<VSBufferReadableStream> {
+	async snapshot(context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream> {
 		if (this.throwOnSnapshot) {
 			throw new Error('Fail');
 		}