Implement shared array buffer file access

jupyterlite · Jun 10, 2024 · 3e2cd9d · 3e2cd9d
1 parent ae234cb
commit 3e2cd9d
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 5 deletions.
diff --git a/packages/pyodide-kernel-extension/src/index.ts b/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,
         });
       },
     });

diff --git a/packages/pyodide-kernel/src/kernel.ts b/packages/pyodide-kernel/src/kernel.ts
@@ -3,13 +3,14 @@ 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,9 +26,31 @@ export class PyodideKernel extends BaseKernel implements IKernel {
     this._worker = this.initWorker(options);
     this._worker.onmessage = (e) => this._processWorkerMessage(e.data);
     this._remoteKernel = coincident(this._worker) as IPyodideWorkerKernel;
+    this._contentsManager = options.contentsManager;
+    this.setupFilesystemAPIs();
     this.initRemote(options);
   }
 
+  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);
+    };
+  }
+
   /**
    * Load the worker.
    *
@@ -288,6 +311,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 +359,10 @@ export namespace PyodideKernel {
       lockFileURL: string;
       packages: string[];
     };
+
+    /**
+     * The Jupyterlite content manager
+     */
+    contentsManager: Contents.IManager;
   }
 }
diff --git a/packages/pyodide-kernel/src/tokens.ts b/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>;
 }
 
 /**

diff --git a/packages/pyodide-kernel/src/worker.ts b/packages/pyodide-kernel/src/worker.ts
@@ -3,10 +3,44 @@
 
 import type Pyodide from 'pyodide';
 
-import type { DriveFS } from '@jupyterlite/contents';
+import { ContentsAPI, DriveFS, ServiceWorkerContentsAPI, TDriveMethod, TDriveRequest, TDriveResponse } from '@jupyterlite/contents';
 
 import type { IPyodideWorkerKernel } from './tokens';
 
+
+/**
+ * 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 PyodideRemoteKernel {
   constructor() {
     this._initialized = new Promise((resolve, reject) => {
@@ -132,9 +166,8 @@ export class PyodideRemoteKernel {
       const mountpoint = '/drive';
       const { FS, PATH, ERRNO_CODES } = this._pyodide;
       const { baseUrl } = options;
-      const { DriveFS } = await import('@jupyterlite/contents');
 
-      const driveFS = new DriveFS({
+      const driveFS = new PyodideDriveFS({
         FS,
         PATH,
         ERRNO_CODES,
@@ -505,5 +538,5 @@ export class PyodideRemoteKernel {
   protected _stdout_stream: any;
   protected _stderr_stream: any;
   protected _resolveInputReply: any;
-  protected _driveFS: DriveFS | null = null;
+  protected _driveFS: PyodideDriveFS | null = null;
 }