Astro + Web Assembly

Let’s figure out how to get web assembly code to bundle and run in the browser in a website built with Astro.

Contents

Web Assembly

First we need some web assembly. Rust has the wasm32-unknown-unknown target that has tier 2 support. This means things will be “guaranteed to build” so it seems like a good language choice. Also I’m bias because rust is great :) Let’s follow the Rust and WebAssembly book until we have the basics in place.

Set up

Install rust toolchain & cargo using rustup.

Install the dependencies we need to build & package rust into wasm:

$ curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
$ cargo install cargo-generate

Generate Project

You can invent your own folder structures here but I created a lib directory in the root of my astro project.

$ mkdir lib && cd lib
$ cargo generate --git https://github.com/rustwasm/wasm-pack-template

I called my project cool-owl-greeting. Your lib directory should look something like this:

$ tree
.
└── cool-owl-greeting
    ├── Cargo.toml
    ├── LICENSE_APACHE
    ├── LICENSE_MIT
    ├── README.md
    ├── src
    │   ├── lib.rs
    │   └── utils.rs
    └── tests
        └── web.rs

Let’s edit lib.rs to prove we can change things and it all works.

#[wasm_bindgen]
pub fn greet() {
    alert("hoot hoot");
}

Compile

Compiling our new project will download the wasm32 rust compilation target, compile to it, and bundle our wasm into the pkg directory.

$ cd lib/cool-owl-greeting
$ wasm-pack build --release
$ tree pkg/
pkg/
├── cool_owl_greeting_bg.js
├── cool_owl_greeting_bg.wasm
├── cool_owl_greeting_bg.wasm.d.ts
├── cool_owl_greeting.d.ts
├── cool_owl_greeting.js
├── package.json
└── README.md

Astro Integration

Let’s think about the kind of API I’d like to have in my markdown blog files in Astro. It would be great if I could render a custom react component and for it to import the cool_owl_greeting.js from our pkg directory in lib, and for it to call into the greet function that is exported.

Rendering React Components in Markdown

From the astro docs we need to install support for mdx. Mdx is a file extension that expands on the functionality of md with features like javascript expressions and components. This is as simple as npx astro add mdx and confirming the astro.config.mjs changes are fine. You should end with the mdx extension in your astro config:

import mdx from "@astrojs/mdx";

export default defineConfig({
  integrations: [
    // ...
    mdx(),
  ],
  // ...
});

Rename the markdown from .md to .mdx. Now we can import and render react components inside our blog posts.

React Component

We can import from the lib directory by adding the path in tsconfig.js:

{
  "extends": "astro/tsconfigs/strict",
  "compilerOptions": {
    "baseUrl": ".",
    "paths": {
      "@lib/*": ["lib/*"], // <-- here
      "@*": ["./src/*"]
    }
  }
}

And import the javascript generated with wasm-pack in our new component:

// CoolOwlGreet.tsx

import { greet } from "@lib/cool-owl-greeting/pkg/cool_owl_greeting";

export default function CoolOwlGreet() {
  return (
    <>
      click <div onClick={() => greet()}>here</div>
    </>
  );
}

This causes vite to error. Line 1 of cool_owl_greeting.js imports the wasm file:

import * as wasm from "./cool_owl_greeting_bg.wasm";

And vite isn’t happy about that:

error   "ESM integration proposal for Wasm" is not supported currently. Use vite-plugin-wasm or other community plugins to handle this. Alter
natively, you can use `.wasm?init` or `.wasm?url`. See https://vitejs.dev/guide/features.html#webassembly for more details.

We can fix this by using the vite-plugin-wasm package and adding it as a plugin in astro.config.mjs:

import wasm from "vite-plugin-wasm";

export default defineConfig({
  // ...
  vite: {
    // ...
    plugins: [wasm()],
  },
});

Now render this component in our mdx file by importing it and including it in the DOM in the same way we would in normal react:

// blog-post.mdx

import CoolOwlGreet from "@components/coolOwlGreeting/CoolOwlGreet";

<CoolOwlGreet client:only="react" />

client:only="react" is a “template directive” and tells astro that the react component should not be server side rendered. Our wasm is imported at runtime and therefore we also need to render at runtime. Read more here

You should see an alert containing “hoot hoot” if you click

Blocked UI Thread

We have a problem that is hiding itself with this solution so far. If we do anything that takes a bit of time in this wasm function then it blocks the UI thread - our browser becomes unresponsive until the function completes. Let’s reproduce the issue.

Repro Steps

We need a new wasm function that takes a long time to return:

#[wasm_bindgen(js_name = "doWork")]
pub fn do_work(delay_ms: f64) -> js_sys::Promise {
    wasm_bindgen_futures::future_to_promise(async move {
        let start = js_sys::Date::now();
        loop {
            let current = js_sys::Date::now();
            if current - start > delay_ms {
                break;
            }
        }
        Ok(JsValue::null())
    })
}

A little to unpack here:

Update our CoolOwlGreet component to call doWork before greet.

// CoolOwlGreet.tsx

import { greet, doWork } from "@lib/cool-owl-greeting/pkg/cool_owl_greeting";

export default function CoolOwlGreet() {
  return <div onClick={() => doWork(1000).then(greet)}>Here</div>;
}

Try to select some text while this is running:

It doesn’t work until we dismiss the alert. Which is good - we can reproduce the issue. We need to be able to free up the javascript UI thread while this runs and - as we’ve seen - simply placing this in a promise does not fix the issue.

Web Workers

We need to put the long running function call in its own thread. In javascript land a thread is a web worker. We’ll also use comlink to make communication between the main thread and web worker easier.

Create a new file worker.ts next to the react component:

// worker.ts

import * as Comlink from "comlink";
import { doWork as wasmDoWork } from "../../../lib/cool-owl-greeting/pkg/cool_owl_greeting";

export class CoolOwl {
  doWork(delay_ms: number) {
    wasmDoWork(delay_ms);
  }
}

Comlink.expose(CoolOwl);

And update our CoolOwlGreet component to use this new worker:

// CoolOwlGreet.tsx

import * as Comlink from "comlink";
import type { CoolOwl } from "./worker";
import { greet } from "@lib/cool-owl-greeting/pkg/cool_owl_greeting";

export default function CoolOwlGreet() {
  const CoolOwlWorker = Comlink.wrap(
    new Worker(new URL("./worker.ts", import.meta.url), {
      type: "module",
    })
  );

  const onClick = useCallback(async () => {
    const instance = await new CoolOwlWorker();
    await instance.doWork(1000).then(greet);
  });

  return (
    <>
      click <div onClick={onClick}>here</div>
    </>
  );
}

We need to massage our vite configuration again. Add wasm and topLevelAwait to the worker plugins:

import topLevelAwait from "vite-plugin-top-level-await";
import wasm from "vite-plugin-wasm";

export default defineConfig({
  // ...
  vite: {
    // ...
    worker: {
      plugins: [wasm(), topLevelAwait()],
    },
  },
});

Try to select some text while this is running:

It works!

Re-Rendering

We have another problem that isn’t showing itself yet. If CoolOwlGreet rerenders for any reason then we re-create our web worker. This imports our worker.js and .wasm files which results in several network requests.

Repro Steps

When do react components re-render? When some state changes. So let’s add some state and get our component to re-render:

export default function CoolOwlGreet() {
  const CoolOwlWorker = Comlink.wrap<typeof CoolOwl>(
    new Worker(new URL("./worker.ts", import.meta.url), {
      type: "module",
    })
  );

  // Dummy value & callback to cause a re-render on click
  const [value, setValue] = useState(0);
  const onClick = useCallback(() => {}, [value, setValue]);

  return (
    <>
      <span onClick={onClick}>here</span>
      <span>{` (counter ${value})`}</span>
    </>
  );
}

Click and the updated counter causes a re-render, causing our web worker to be reloaded, causing (albeit cached) network requests. If you open the network inspector you can see it in action:

network-inspector

Great, again we can repro, so let’s fix.

cool owl says:

Why fix if they are already cached, what’s the issue?

I suppose there is no issue… Functionality wise we are totally fine. But it’s messy, inconsistent. It’s a side effect of our code that we didn’t expect or want. It’s the browser picking up the pieces after our sloppy coding, and we can do better.

Memoisation

We can use memoisation to solve this. React has useMemo so let’s use that:

import { useMemo } from "react";

export default function CoolOwlGreet() {
  const CoolOwlWorker = useMemo(
    () =>
      Comlink.wrap<typeof CoolOwl>(
        new Worker(new URL("./worker.ts", import.meta.url), {
          type: "module",
        })
      ),
    [] // <-- no dependencies
  );

  // ...
}

We pass an empty array to the dependencies, we never need to reload the web worker.

Click and see that we have fixed things.

Conclusion

We learned how to compile Rust to wasm, and how to import that into a blog post in astro. We also learned how to not block the UI thread while our wasm functions are running and how to now reload our wasm code on component re-renders. This opens up a lot of opportunity for us for cool things…