allow returning undefined if file was not saved (#213123)

* allow returning undefined if file was not saved

* bring back cancellation check

* helper utility to convey optional result if cancelled

* edit

* inline function

* test cancelled custom save

* inline more

---------

Co-authored-by: Benjamin Pasero <benjamin.pasero@microsoft.com>

src/vs/workbench/api/common/extHostNotebook.ts

@@ -335,7 +335,6 @@ export class ExtHostNotebookController implements ExtHostNotebookShape {
 			throw new files.FileOperationError(localize('err.readonly', "Unable to modify read-only file '{0}'", this._resourceForError(uri)), files.FileOperationResult.FILE_PERMISSION_DENIED);
 		}
 
-
 		const data: vscode.NotebookData = {
 			metadata: filter(document.apiNotebook.metadata, key => !(serializer.options?.transientDocumentMetadata ?? {})[key]),
 			cells: [],
@@ -360,7 +359,15 @@ export class ExtHostNotebookController implements ExtHostNotebookShape {
 		// validate write
 		await this._validateWriteFile(uri, options);
 
+		if (token.isCancellationRequested) {
+			throw new Error('canceled');
+		}
 		const bytes = await serializer.serializer.serializeNotebook(data, token);
+		if (token.isCancellationRequested) {
+			throw new Error('canceled');
+		}
+
+		// Don't accept any cancellation beyond this point, we need to report the result of the file write
 		this.trace(`serialized versionId: ${versionId} ${uri.toString()}`);
 		await this._extHostFileSystem.value.writeFile(uri, bytes);
 		this.trace(`Finished write versionId: ${versionId} ${uri.toString()}`);

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

@@ -35,16 +35,6 @@ import { IMarkdownString } from 'vs/base/common/htmlContent';
  */
 export interface IStoredFileWorkingCopyModelFactory<M extends IStoredFileWorkingCopyModel> extends IFileWorkingCopyModelFactory<M> { }
 
-export async function createOptionalResult<T>(callback: (token: CancellationToken) => Promise<T | undefined>, token: CancellationToken): Promise<T | undefined> {
-	const result = await callback(token);
-	if (result === undefined && token.isCancellationRequested) {
-		return undefined;
-	}
-	else {
-		return assertIsDefined(result);
-	}
-}
-
 /**
  * The underlying model of a stored file working copy provides some
  * methods for the stored file working copy to function. The model is
@@ -1029,7 +1019,15 @@ export class StoredFileWorkingCopy<M extends IStoredFileWorkingCopyModel> extend
 
 					// Delegate to working copy model save method if any
 					if (typeof resolvedFileWorkingCopy.model.save === 'function') {
-						stat = await resolvedFileWorkingCopy.model.save(writeFileOptions, saveCancellation.token);
+						try {
+							stat = await resolvedFileWorkingCopy.model.save(writeFileOptions, saveCancellation.token);
+						} catch (error) {
+							if (saveCancellation.token.isCancellationRequested) {
+								return undefined; // save was cancelled
+							}
+
+							throw error;
+						}
 					}
 
 					// Otherwise ask for a snapshot and save via file services

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

@@ -88,12 +88,21 @@ export class TestStoredFileWorkingCopyModelWithCustomSave extends TestStoredFile
 
 	saveCounter = 0;
 	throwOnSave = false;
+	saveOperation: Promise<void> | undefined = undefined;
 
 	async save(options: IWriteFileOptions, token: CancellationToken): Promise<IFileStatWithMetadata> {
 		if (this.throwOnSave) {
 			throw new Error('Fail');
 		}
 
+		if (this.saveOperation) {
+			await this.saveOperation;
+		}
+
+		if (token.isCancellationRequested) {
+			throw new Error('Canceled');
+		}
+
 		this.saveCounter++;
 
 		return {
@@ -190,6 +199,42 @@ suite('StoredFileWorkingCopy (with custom save)', function () {
 		assert.strictEqual(workingCopy.hasState(StoredFileWorkingCopyState.ERROR), true);
 	});
 
+	test('save cancelled (custom implemented)', async () => {
+		let savedCounter = 0;
+		let lastSaveEvent: IStoredFileWorkingCopySaveEvent | undefined = undefined;
+		disposables.add(workingCopy.onDidSave(e => {
+			savedCounter++;
+			lastSaveEvent = e;
+		}));
+
+		let saveErrorCounter = 0;
+		disposables.add(workingCopy.onDidSaveError(() => {
+			saveErrorCounter++;
+		}));
+
+		await workingCopy.resolve();
+		let resolve: () => void;
+		(workingCopy.model as TestStoredFileWorkingCopyModelWithCustomSave).saveOperation = new Promise(r => resolve = r);
+
+		workingCopy.model?.updateContents('first');
+		const firstSave = workingCopy.save();
+		// cancel the first save by requesting a second while it is still mid operation
+		workingCopy.model?.updateContents('second');
+		const secondSave = workingCopy.save();
+		resolve!();
+		await firstSave;
+		await secondSave;
+
+		assert.strictEqual(savedCounter, 1);
+		assert.strictEqual(saveErrorCounter, 0);
+		assert.strictEqual(workingCopy.isDirty(), false);
+		assert.strictEqual(lastSaveEvent!.reason, SaveReason.EXPLICIT);
+		assert.ok(lastSaveEvent!.stat);
+		assert.ok(isStoredFileWorkingCopySaveEvent(lastSaveEvent!));
+		assert.strictEqual(workingCopy.model?.pushedStackElement, true);
+		assert.strictEqual((workingCopy.model as TestStoredFileWorkingCopyModelWithCustomSave).saveCounter, 1);
+	});
+
 	ensureNoDisposablesAreLeakedInTestSuite();
 });