FileSystem calls over Atomics.wait instead of service worker when ava…

…ilable (#114)

* Implement shared array buffer file access

* Linter

* Iterate

* Reduce diff

* Debug

packages/pyodide-kernel-extension/src/index.ts

@@ -41,6 +41,8 @@ const kernel: JupyterLiteServerPlugin<void> = {
     serviceWorker?: IServiceWorkerManager,
     broadcastChannel?: IBroadcastChannelWrapper,
   ) => {
+    const contentsManager = app.serviceManager.contents;
+
     const config =
       JSON.parse(PageConfig.getOption('litePluginSettings') || '{}')[PLUGIN_ID] || {};
 
@@ -93,6 +95,7 @@ const kernel: JupyterLiteServerPlugin<void> = {
           disablePyPIFallback,
           mountDrive,
           loadPyodideOptions,
+          contentsManager,
         });
       },
     });

packages/pyodide-kernel/package.json

@@ -37,7 +37,7 @@
     "build:lib": "tsc -b",
     "build:prod": "jlpm build",
     "build:py": "python scripts/generate-wheels-js.py",
-    "build:worker": "esbuild --bundle --minify --sourcemap --target=es2019 --format=esm --outfile=lib/worker.js src/worker.ts",
+    "build:worker": "esbuild --bundle --minify --sourcemap --target=es2019 --format=esm --outfile=lib/coincident.worker.js src/coincident.worker.ts",
     "dist": "cd ../../dist && npm pack ../packages/pyodide-kernel",
     "clean": "jlpm clean:lib && jlpm clean:py",
     "clean:all": "jlpm clean",

packages/pyodide-kernel/src/coincident.worker.ts

@@ -6,13 +6,83 @@
  */
 import coincident from 'coincident';
 
+import {
+  ContentsAPI,
+  DriveFS,
+  ServiceWorkerContentsAPI,
+  TDriveMethod,
+  TDriveRequest,
+  TDriveResponse,
+} from '@jupyterlite/contents';
+
 import { PyodideRemoteKernel } from './worker';
 import { IPyodideWorkerKernel } from './tokens';
 
-const worker = new PyodideRemoteKernel();
-
 const workerAPI: IPyodideWorkerKernel = coincident(self) as IPyodideWorkerKernel;
 
+/**
+ * An Emscripten-compatible synchronous Contents API using shared array buffers.
+ */
+export class SharedBufferContentsAPI extends ContentsAPI {
+  request<T extends TDriveMethod>(data: TDriveRequest<T>): TDriveResponse<T> {
+    return workerAPI.processDriveRequest(data);
+  }
+}
+
+/**
+ * A custom drive implementation which uses shared array buffers if available, service worker otherwise
+ */
+class PyodideDriveFS extends DriveFS {
+  createAPI(options: DriveFS.IOptions): ContentsAPI {
+    if (crossOriginIsolated) {
+      return new SharedBufferContentsAPI(
+        options.driveName,
+        options.mountpoint,
+        options.FS,
+        options.ERRNO_CODES,
+      );
+    } else {
+      return new ServiceWorkerContentsAPI(
+        options.baseUrl,
+        options.driveName,
+        options.mountpoint,
+        options.FS,
+        options.ERRNO_CODES,
+      );
+    }
+  }
+}
+
+export class PyodideCoincidentKernel extends PyodideRemoteKernel {
+  /**
+   * Setup custom Emscripten FileSystem
+   */
+  protected async initFilesystem(
+    options: IPyodideWorkerKernel.IOptions,
+  ): Promise<void> {
+    if (options.mountDrive) {
+      const mountpoint = '/drive';
+      const { FS, PATH, ERRNO_CODES } = this._pyodide;
+      const { baseUrl } = options;
+
+      const driveFS = new PyodideDriveFS({
+        FS,
+        PATH,
+        ERRNO_CODES,
+        baseUrl,
+        driveName: this._driveName,
+        mountpoint,
+      });
+      FS.mkdir(mountpoint);
+      FS.mount(driveFS, {}, mountpoint);
+      FS.chdir(mountpoint);
+      this._driveFS = driveFS;
+    }
+  }
+}
+
+const worker = new PyodideCoincidentKernel();
+
 workerAPI.initialize = worker.initialize.bind(worker);
 workerAPI.execute = worker.execute.bind(worker);
 workerAPI.complete = worker.complete.bind(worker);

packages/pyodide-kernel/src/kernel.ts

@@ -3,13 +3,18 @@ import coincident from 'coincident';
 import { PromiseDelegate } from '@lumino/coreutils';
 
 import { PageConfig } from '@jupyterlab/coreutils';
-import { KernelMessage } from '@jupyterlab/services';
+import { Contents, KernelMessage } from '@jupyterlab/services';
 
 import { BaseKernel, IKernel } from '@jupyterlite/kernel';
 
 import { IPyodideWorkerKernel, IRemotePyodideWorkerKernel } from './tokens';
 
 import { allJSONUrl, pipliteWheelUrl } from './_pypi';
+import {
+  DriveContentsProcessor,
+  TDriveMethod,
+  TDriveRequest,
+} from '@jupyterlite/contents';
 
 /**
  * A kernel that executes Python code with Pyodide.
@@ -25,6 +30,28 @@ export class PyodideKernel extends BaseKernel implements IKernel {
     this._worker = this.initWorker(options);
     this._worker.onmessage = (e) => this._processWorkerMessage(e.data);
     this._remoteKernel = this.initRemote(options);
+    this._contentsManager = options.contentsManager;
+    this.setupFilesystemAPIs();
+  }
+
+  private setupFilesystemAPIs() {
+    (this._remoteKernel.processDriveRequest as any) = async <T extends TDriveMethod>(
+      data: TDriveRequest<T>,
+    ) => {
+      if (!DriveContentsProcessor) {
+        throw new Error(
+          'File system calls over Atomics.wait is only supported with jupyterlite>=0.4.0a3',
+        );
+      }
+
+      if (this._contentsProcessor === undefined) {
+        this._contentsProcessor = new DriveContentsProcessor({
+          contentsManager: this._contentsManager,
+        });
+      }
+
+      return await this._contentsProcessor.processDriveRequest(data);
+    };
   }
 
   /**
@@ -288,6 +315,8 @@ export class PyodideKernel extends BaseKernel implements IKernel {
     return await this._remoteKernel.inputReply(content, this.parent);
   }
 
+  private _contentsManager: Contents.IManager;
+  private _contentsProcessor: DriveContentsProcessor | undefined;
   private _worker: Worker;
   private _remoteKernel: IRemotePyodideWorkerKernel;
   private _ready = new PromiseDelegate<void>();
@@ -334,5 +363,10 @@ export namespace PyodideKernel {
       lockFileURL: string;
       packages: string[];
     };
+
+    /**
+     * The Jupyterlite content manager
+     */
+    contentsManager: Contents.IManager;
   }
 }

packages/pyodide-kernel/src/tokens.ts

@@ -5,6 +5,7 @@
  * Definitions for the Pyodide kernel.
  */
 
+import { TDriveMethod, TDriveRequest, TDriveResponse } from '@jupyterlite/contents';
 import { IWorkerKernel } from '@jupyterlite/kernel';
 
 /**
@@ -20,6 +21,14 @@ export interface IPyodideWorkerKernel extends IWorkerKernel {
    * Handle any lazy initialization activities.
    */
   initialize(options: IPyodideWorkerKernel.IOptions): Promise<void>;
+
+  /**
+   * Process drive request
+   * @param data
+   */
+  processDriveRequest<T extends TDriveMethod>(
+    data: TDriveRequest<T>,
+  ): TDriveResponse<T>;
 }
 
 /**