Swappable WASIDrive Implementations

[jacobrobertsbaca]

Hi there!

I am interested in exploring the discussion at the end of #244 more. Right now, there is not a good way to interface with file reads/writes on the JS end: we can use the hack in #244 to create fake Uint8Array objects and pass them as the content of a WASIFile, but it would be better to have a more generic interface built into runno itself.

I propose adding a swappable WASIDrive implementation to WASIContext, e.g.:

export class WASIContext {
  drive: WASIDrive;
  /* ... */
}

and then implementing drive creation in WASI in terms of context.drive. This would allow creating custom WASIDrive implementations that might, e.g. support observing file reads:

class CustomDrive extends WASIDrive {
  onFileRead: (fd: FileDescriptor) => void;

  read(fd: FileDescriptor, bytes: number): DriveResult<Uint8Array> {
    this.onFileRead(fd);
    return super.read(fd, bytes);
  }
}

WASI.start(/* ... */, { drive: new CustomDrive(), /* ... */ });

However, this interface leaves something to be desired as it seems like there's some undesirable overlap between this new drive option and the existing fs in WASIContext. For this, I propose two options.

Option A: Breaking Changes

In the extreme case, an implementation could entirely obviate the use of a WASIFS (a custom WASIDrive might choose not to use a backing WASIFS at all), so one option is to have an interface like:

export abstract class FileDrive {
  // Define abstract versions of all system calls
  abstract read(fd: FileDescriptor, bytes: number): DriveResult<Uint8Array>;
  // ...
};

export class WASIDrive extends FileDrive {
  constructor(fs: WASIFS) { /* ... */ };
  /* Current implementation of WASIDrive... */
};

and then for backwards compatibility, have the WASIContext look like:

export class WASIContext {
  fs: WASIFS | FileDrive;
  /* ... */
}

If a WASIFS is passed, a WASIDrive gets created from instead.

However, this probably goes a step too far, since the WASIExecutionResult is built around a WASIFS--how do we get the corresponding fs from the FileDrive? We could change WASIExecutionResult to return a FileDrive as its result instead, but this would likely lead to way too many breaking changes as it fundamentally changes the @runno/wasi interface.

Option B: Preferred Version

A more conservative option with no or fewer breaking changes is to keep the WASIDrive tethered around a WASIFS, at least in terms of interface, perhaps like so:

export abstract class FileDrive {

  // Every FileDrive is built around a WASIFS
  abstract get fs(): WASIFS;
  constructor(fs: WASIFS) {}

  // Define abstract versions of all system calls
  abstract read(fd: FileDescriptor, bytes: number): DriveResult<Uint8Array>;
  // ...
};

export class WASIDrive extends FileDrive {
  get fs(): WASIFS { return this.fs; }
  constructor(public fs: WASIFS) {
    super(fs);
    /* ... */
  };

  /* Current implementation of WASIDrive... */
};

and then define:

export class WASIContext {
  fs: WASIFS | FileDrive;
  /* ... */
}

Again, such that for the fs we can either pass a WASIFS, from which a WASIDrive gets created, or our own FileDrive. If a custom FileDrive implementation does not need to use a WASIFS for whatever reason, they could just pass {} to the FileDrive super constructor. This also does not lead to any breaking changes to WASIExecutionResult, since it just returns drive.fs for the result fs.

Use Case: Interactive Canvas

My primary use case for this is to enable interaction with a WASIWorkerHost beyond stdin/stdout. In my particular case, I am looking into creating an interactive canvas library driven through the filesystem. An executing runno program could write to a special file on the filesystem to draw something to the screen, and my custom FileDrive implementation could observe that write and render to the canvas instead. For example:

class CanvasFileDrive extends WASIDrive {
  write(fd: FileDescriptor, data: Uint8Array): Result {
    if (isCanvasFd(fd)) return writeToCanvas(data);
    else return super.write(fd, data);
  }

  private writeToCanvas(fd: FileDescriptor, data: Uint8Array): Result {
    // Decodes data and draws to the canvas
  }
};

However, this system would also enable receiving data as well inside runno--if we wanted a runno program to read something from the browser, e.g. get the current mouse position, we could do something similar where we override the read system call and the browser writes the current coordinates.

Overview

This is effectively a way of adding imports/exports to the WebAssembly module, but keeps the spirit of the runno interface. It also lines up well with the idea of a Unix synthetic filesystem where special entries (e.g. /dev/null) do not actually correspond to physical files (in this case the WASIFS) but special behaviour to be implemented by the OS (e.g. us).

One issue #244 presents is how this interface would handle a backing filesystem whose operations are fundamentally asynchronous, e.g. the fact that file handle creation in OPFS is asynchronous. I am not sure how to handle this aside from making the FileDrive operations themselves asynchronous, but I don't know if that would play nicely with the wasm exports. This might just be out of scope for this change, and it could just be easier for now to mandate that FileDrive implementations are synchronous.

Would love to hear some thoughts on fine-tuning this interface/thinking through any gaps/breaking changes I may have missed. I am happy to submit a PR once there are some thoughts on what a good interface for this looks like.

[taybenlor]

Okay these are great thoughts. Thanks heaps for taking the time to write it out.

An important consideration will be the difference between running WASI directly with WASI.start, and running it through the WASIWorkerHost. In the first context we can easily pass complicated objects. In the second case we have to be careful about the boundary, as we need to communicate via postMessage (and in the rare case SharedArrayBuffer). We can quite easily pass UInt8Array and string objects in that context (see: Transferable Objects on MDN). I'd recommend reading the implementation in wasi-host.ts and wasi-worker.ts.

I'd also like to (where possible) maintain conceptual compatibility with the way Node's WASI API works (docs). They have no way of representing something like this, we'll just need to make sure there are no overlaps in options that are incompatible.

I am not sure how to handle this aside from making the FileDrive operations themselves asynchronous

You're right, when running WASI within the current thread all operations must return immediately. However if it is run in a Worker then we can block on the Worker side and wait for a response, this is how Runno is able to do interactive stdin/stdout.

Some ideas for how I've been imagining this working follow.

First Idea - Simplified Abstraction through Callbacks

When I was considering this I was imagining some kind of object based interface with a few special handler functions that create a simplified abstraction over a filesystem:

WASI.start(/* ... */, {
  customFS: {
    exists: (path: string): boolean => {},
    read: (path: string): UInt8Array => {},
    write: (path: string, content: UInt8Array) => {},
    
    // Note: We will also need to consider how this works for directories
  },
  /* ... */
});

The concept would be to abstract away all the details of reading particular ranges of bytes, writing particular ranges etc. There would be an implementation on the @runno/wasi side that continues to manage WASIDrive. The entire file contents would be sent for each read and write. This would offer lower performance but much simpler cognitive complexity. Since it's all happening in memory it should be relatively fast anyway.

Second Idea - Copy the OPFS API

Instead of inventing our own API, we could take a FileSystemDirectoryHandle as an argument (mdn docs). This could literally be a handle from OPFS, allowing for reading/writing to the real file-system. Or it could be the user's own implementation of the same API.

We could provide our own abstraction like RunnoFileSystemDirectoryHandle that takes arguments similar to the callback option above to give users the ability to run their own fully customised file system.

Internally we would still need to manage a synchronisation process to/from the WASIWorker.

Summary

Some considerations that will need to be made:

1. The WASIWorkerHost needs to communicate to the worker via postMessage to work effectively.
2. API consumers will want a conceptually simple API or they will bounce off it
3. If we can achieve OPFS compatibility using this, that would be handy
4. The internal boundary between WASIWorkerHost and WASIWorker will constrain implementation

I'd love to hear your thoughts!

[jacobrobertsbaca]

Hey @taybenlor, thanks for the response. It seems I did not fully appreciate the difficulty of passing data to the web worker, however this may also be an opportunity to integrate async-backed file systems through the use of Atomics as you mentioned. I think we can achieve both the simplicity of the callbacks approach together with the flexibility of supporting something like OPFS by extending the existing WASIDrive with new implementations. The benefit of this approach is that it would assume nothing about the user's filesystem: they could use one of the pre-built ones for common backends (e.g. OPFS) or swap it out with their own. Here's what I'm thinking.

Synchronous Filesystems

If we are running a WASI instance outside of a web worker, all filesystem operations will need to be synchronous since the WASI Preview 1 exports are synchronous. We can define a common interface for the synchronous system calls:

export interface SyncDrive {
  fs: WASIFS;
  open(fdDir: FileDescriptor, path: WASIPath, oflags: number, fdflags: number): DriveResult<FileDescriptor>;
  close(fd: FileDescriptor): Result;
  read(fd: FileDescriptor, bytes: number): DriveResult<Uint8Array>;
  write(fd: FileDescriptor, data: Uint8Array): Result;

  /* Other synchronous system call definitions */
}

export class WASIDrive implements SyncDrive {
  constructor(public fs: WASIFS) { /* ... */ }

  /* Existing WASIDrive implementations of system calls */
}

and, as mentioned, we can extend the WASIContext.fs to support SyncDrive:

export class WASIContext {
  fs: WASIFS | SyncDrive;
  /* ... */
}

As you stated, this is a bit much for most consumers, but we could simplify the API a bit by providing, e.g. a CallbacksDrive implemented through user-passed callbacks:

export class CallbacksDrive implements SyncDrive { /* ... */ }

WASI.start(/* ... */, {
  fs: new CallbacksDrive({
    exists: (path: string): boolean => {},
    read: (path: string): UInt8Array => {},
    write: (path: string, content: UInt8Array) => {},
    /* ... */
  })
});

Asynchronous Filesystems

As you noted, we can't just pass the SyncDrive object to the web worker since it's not transferable. We could, however, keep the drive parked on the browser main thread and, within the worker, create a SyncDrive that communicates with the main thread via postMessage, SharedArrayBuffer, and Atomics to a (possibly asynchronous) filesystem. In effect, the WASIWorkerHost would provide the glue to map a SyncDrive (passed to WASI.start, implemented using Atomics) into an asynchronous filesystem running on the main thread. This would allow suppport for asynchronous backends like OPFS, and also manipulating browser resources (e.g. Canvas) without any special logic to transfer them to the web worker.

export type AsyncDrive = {
  [K in keyof SyncDrive]: SyncDrive[K] extends (...args: infer A) => infer R
    ? (...args: A) => R | Promise<R>
    : SyncDrive[K];
};

type WASIWorkerHostContext = Partial<Omit<WASIContextOptions, "stdin" | "fs">> & {
  fs: WASIFS | AsyncDrive;
};

export class WASIWorkerHost {
  constructor(binaryURL: string, context: WASIWorkerHostContext) { /* ... */ }
};

Notice that every SyncDrive is also an AsyncDrive. Every drive we could pass to WASI we can also pass to WASIWorkerHost, but we could also create asynchronous implementations, for example for OPFS:

export class OPFSDrive implements AsyncDrive {
  readonly fs: WASIFS = {}; // Unused
  root: FileSystemDirectoryHandle; // Opened in constructor

  async open(fdDir: FileDescriptor, path: WASIPath, oflags: number, fdflags: number): Promise<DriveResult<FileDescriptor>> {
    const file = await this.root.getFileHandle(path);
    // Rest of implementation...
  }
}

Implementation

Without going into too many specifics, here's what would happen during a read system call made by a WASI binary running inside a WASIWorkerHost:

1. Binary calls read import
2. WASI instance inside WASIWorker calls corresponding read method on a custom SyncDrive implementation
3. SyncDrive calls postMessage with information about which method was called and the serialized arguments.
4. SyncDrive uses Atomics.wait on a shared buffer to await the response synchronously.
5. WASIWorkerHost receives the read request in its onmessage callback.
6. WASIWorkerHost calls read on its AsyncDrive
7. When the read promise resolves, it writes the (serialized) result into the shared buffer and calls Atomics.notify to wake the web worker.
8. SyncDrive deserializes the result from its shared buffer upon waking and returns it.
9. Binary consumes returned bytes.

---

To address your considerations, I think:

1. If WASIWorker is waiting to be woken up after a call to Atomics.wait, I don't think it will be able to respond to a postMessage. For that reason, we can write to a shared buffer any time we need to pass data from the browser thread to the web worker and then use Atomics.notify to wake up the worker.
2. Agree on simplicity of the API: we can achieve this through subclassing and providing a CallbacksDrive (or something along these lines). At the same time, if someone wants to roll their own filesystem, they should have the full power to do that however they'd like and should not be constrained by a simplified interface!
3. I think we can do this through the AsyncDrive. That said, I don't think it's worthwhile to lock users in to this particular implementation--what if someone wants to bypass OPFS altogether? For instance, what if their filesystem is stored remotely on a database? Extending the existing WASIDrive gives users the power to implement their own filesystem with its own peculiarities.
4. This boundary exists, but should be close to non-existent with the above approach.

Let me know what you think! I am still familiarizing myself with web workers, so if I said something that sounds impossible, please let me know! I am a bit concerned with serialization/message passing overhead with this approach and the performance of blocking with Atomics, so would be curious if you think this would be a concern.

[taybenlor]

I love this! 😍

Thankyou so much for going deep on how this could work with Web Workers. I particularly like the idea that every SyncDrive is also an AsyncDrive.

I can't see any obvious flaws at this point, though to give you context I'm currently writing this on my phone from a plane.

The next things to think about are how do we tackle this? I tend to like thinking from the user backward. So defining a nice DX (developer experience) for that CallbackDrive will be essential.

Here's my proposed breakdown:

• Define a rough goal for the CallbackDrive DX (we're basically there)
• Write an implementation of the SyncDrive - doing this first will help understand the problem and get some progress that we can ship
• Build the SyncDrive <-> AsyncDrive architecture (I think this will need some more exploring and defining, there are some more sophisticated Atomics primitives in baseline now)
• Build the CallbackDrive and the OPFSDrive but now understanding how it actually should work (it won't work the same way we originally planned because it never does)
• Write a blog post or something!

What do you think?

[jacobrobertsbaca]

Amazing, I think this sounds like a good plan. I think this can sort of be divided into two parts: the easier part is getting the SyncDrive functionality working with the existing WASI interface (but keeping it excluded from the WASIWorkerHost); the harder part is building out the AsyncDrive + support for WASIWorkerHost.

I can start a fork and begin working on an initial prototype for this. As for the interface of the CallbacksDrive, the approach you suggested:

{
    exists: (path: string): boolean => {},
    read: (path: string): UInt8Array => {},
    write: (path: string, content: UInt8Array) => {},
    /* ... */
}

seems to function less as "user is getting a callback for a filesystem operation" and more like "user provides the implementation of each filesystem operation". This could be useful, but it does raise the question--what if a user doesn't provide a particular callback? Should it default to using the existing WASIDrive functionality?