Restore changes to ipynb extension to serialize notebook in node worker for desktop

extensions/ipynb/package.json

@@ -39,6 +39,12 @@
             "scope": "resource",
             "markdownDescription": "%ipynb.pasteImagesAsAttachments.enabled%",
             "default": true
+          },
+          "ipynb.experimental.serialization": {
+            "type": "boolean",
+            "scope": "resource",
+            "markdownDescription": "%ipynb.experimental.serialization%",
+            "default": false
           }
         }
       }

extensions/ipynb/src/ipynbMain.browser.ts

@@ -5,9 +5,10 @@
 
 import * as vscode from 'vscode';
 import * as main from './ipynbMain';
+import { NotebookSerializer } from './notebookSerializer.web';
 
 export function activate(context: vscode.ExtensionContext) {
-	return main.activate(context, true);
+	return main.activate(context, new NotebookSerializer(context));
 }
 
 export function deactivate() {

extensions/ipynb/src/ipynbMain.node.ts

@@ -5,9 +5,10 @@
 
 import * as vscode from 'vscode';
 import * as main from './ipynbMain';
+import { NotebookSerializer } from './notebookSerializer.node';
 
 export function activate(context: vscode.ExtensionContext) {
-	return main.activate(context, false);
+	return main.activate(context, new NotebookSerializer(context));
 }
 
 export function deactivate() {

extensions/ipynb/src/ipynbMain.ts

@@ -4,7 +4,6 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
-import { NotebookSerializer } from './notebookSerializer';
 import { activate as keepNotebookModelStoreInSync } from './notebookModelStoreSync';
 import { notebookImagePasteSetup } from './notebookImagePaste';
 import { AttachmentCleaner } from './notebookAttachmentCleaner';
@@ -29,8 +28,7 @@ type NotebookMetadata = {
 	[propName: string]: unknown;
 };
 
-export function activate(context: vscode.ExtensionContext, isBrowser: boolean) {
-	const serializer = new NotebookSerializer(context, isBrowser);
+export function activate(context: vscode.ExtensionContext, serializer: vscode.NotebookSerializer) {
 	keepNotebookModelStoreInSync(context);
 	context.subscriptions.push(vscode.workspace.registerNotebookSerializer('jupyter-notebook', serializer, {
 		transientOutputs: false,

extensions/ipynb/src/notebookSerializer.node.ts

@@ -0,0 +1,91 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import { DeferredPromise, generateUuid } from './helper';
+import { NotebookSerializerBase } from './notebookSerializer';
+
+export class NotebookSerializer extends NotebookSerializerBase {
+	private experimentalSave = vscode.workspace.getConfiguration('ipynb').get('experimental.serialization', false);
+	private worker?: import('node:worker_threads').Worker;
+	private tasks = new Map<string, DeferredPromise<Uint8Array>>();
+
+	constructor(context: vscode.ExtensionContext) {
+		super(context);
+		context.subscriptions.push(vscode.workspace.onDidChangeConfiguration(e => {
+			if (e.affectsConfiguration('ipynb.experimental.serialization')) {
+				this.experimentalSave = vscode.workspace.getConfiguration('ipynb').get('experimental.serialization', false);
+			}
+		}));
+	}
+
+	override dispose() {
+		try {
+			void this.worker?.terminate();
+		} catch {
+			//
+		}
+		super.dispose();
+	}
+
+	public override async serializeNotebook(data: vscode.NotebookData, token: vscode.CancellationToken): Promise<Uint8Array> {
+		if (this.disposed) {
+			return new Uint8Array(0);
+		}
+
+		if (this.experimentalSave) {
+			return this.serializeViaWorker(data);
+		}
+
+		return super.serializeNotebook(data, token);
+	}
+
+	private async startWorker() {
+		if (this.disposed) {
+			throw new Error('Serializer disposed');
+		}
+		if (this.worker) {
+			return this.worker;
+		}
+		const { Worker } = await import('node:worker_threads');
+		const outputDir = getOutputDir(this.context);
+		this.worker = new Worker(vscode.Uri.joinPath(this.context.extensionUri, outputDir, 'notebookSerializerWorker.js').fsPath, {});
+		this.worker.on('exit', (exitCode) => {
+			if (!this.disposed) {
+				console.error(`IPynb Notebook Serializer Worker exited unexpectedly`, exitCode);
+			}
+			this.worker = undefined;
+		});
+		this.worker.on('message', (result: { data: Uint8Array; id: string }) => {
+			const task = this.tasks.get(result.id);
+			if (task) {
+				task.complete(result.data);
+				this.tasks.delete(result.id);
+			}
+		});
+		this.worker.on('error', (err) => {
+			if (!this.disposed) {
+				console.error(`IPynb Notebook Serializer Worker errored unexpectedly`, err);
+			}
+		});
+		return this.worker;
+	}
+	private async serializeViaWorker(data: vscode.NotebookData): Promise<Uint8Array> {
+		const worker = await this.startWorker();
+		const id = generateUuid();
+
+		const deferred = new DeferredPromise<Uint8Array>();
+		this.tasks.set(id, deferred);
+		worker.postMessage({ data, id });
+
+		return deferred.p;
+	}
+}
+
+
+function getOutputDir(context: vscode.ExtensionContext): string {
+	const main = context.extension.packageJSON.main as string;
+	return main.indexOf('/dist/') !== -1 ? 'dist' : 'out';
+}

extensions/ipynb/src/notebookSerializer.ts

@@ -10,21 +10,14 @@ import { getPreferredLanguage, jupyterNotebookModelToNotebookData } from './dese
 import * as fnv from '@enonic/fnv-plus';
 import { serializeNotebookToString } from './serializers';
 
-export class NotebookSerializer extends vscode.Disposable implements vscode.NotebookSerializer {
-	private disposed: boolean = false;
-	private worker?: import('node:worker_threads').Worker;
-
-	constructor(readonly context: vscode.ExtensionContext, _isBrowser: boolean) {
+export abstract class NotebookSerializerBase extends vscode.Disposable implements vscode.NotebookSerializer {
+	protected disposed: boolean = false;
+	constructor(protected readonly context: vscode.ExtensionContext) {
 		super(() => { });
 	}
 
 	override dispose() {
 		this.disposed = true;
-		try {
-			void this.worker?.terminate();
-		} catch {
-			//
-		}
 		super.dispose();
 	}
 
@@ -91,5 +84,5 @@ export class NotebookSerializer extends vscode.Disposable implements vscode.Note
 		const serialized = serializeNotebookToString(data);
 		return new TextEncoder().encode(serialized);
 	}
-}
 
+}

extensions/ipynb/src/notebookSerializer.web.ts

@@ -0,0 +1,9 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { NotebookSerializerBase } from './notebookSerializer';
+
+export class NotebookSerializer extends NotebookSerializerBase {
+}