Merge pull request #20 from zebp/zebp/benchmarking

chore: add benchmarking suite

test/Makefile

@@ -1,6 +1,8 @@
 default: run-tests
 
 WASM_OPT   := ../build/binaryen/bin/wasm-opt
+WASI_SDK_PATH := ../deps/wasi-sdk-13.0
+WASI_SYSROOT  := $(abspath ${WASI_SDK_PATH}/share/wasi-sysroot)
 OUTPUT_DIR := ../build/test
 
 # https://github.com/caspervonb/wasi-test-suite
@@ -21,6 +23,11 @@ WASMTIME_DST := $(OUTPUT_DIR)/wasmtime
 WASMTIME_SRC_TESTS := $(shell ls $(WASI_TESTS_CRATE_PATH)/src/bin/*.rs)
 WASMTIME_DST_TESTS := $(WASMTIME_SRC_TESTS:$(WASMTIME_SRC)/%.rs=$(WASMTIME_DST)/%.wasm)
 
+BENCHMARK_SRC := ./subjects
+BENCHMARK_DST := $(OUTPUT_DIR)/benchmark
+BENCHMARK_SRC_TESTS := $(shell ls ./subjects/*.c)
+BENCHMARK_DST_TESTS := $(BENCHMARK_SRC_TESTS:$(BENCHMARK_SRC)/%.c=$(BENCHMARK_DST)/%.wasm)
+
 $(WASMTIME_BIN)/%.wasm: $(wildcard $(WASI_TESTS_CRATE_PATH)/src/**)
 	cargo build --bin $* --target wasm32-wasi --manifest-path $(WASI_TESTS_CRATE_PATH)/Cargo.toml
 
@@ -36,9 +43,12 @@ node_modules: ./package.json ./package-lock.json ../package.json
 	npm install --no-audit --no-fund --no-progress --quiet
 	touch $@
 
-run-tests: $(BUNDLE) $(WASI_TEST_SUITE_DST_TESTS) $(WASMTIME_DST_TESTS) $(OUTPUT_DIR)/wasm-table.ts node_modules
+$(OUTPUT_DIR)/standalone.mjs: $(BENCHMARK_DST_TESTS) node_modules ./driver/*.ts
+	node benchmark-build.mjs
+
+run-tests: $(BUNDLE) $(OUTPUT_DIR)/wasm-table.ts node_modules $(OUTPUT_DIR)/standalone.mjs
 	$(shell npm bin)/tsc -p ./tsconfig.json
-	JEST_JUNIT_OUTPUT_DIR=$(OUTPUT_DIR) NODE_OPTIONS=--experimental-vm-modules NODE_NO_WARNINGS=1 \
+	JEST_JUNIT_OUTPUT_DIR=$(OUTPUT_DIR) OUTPUT_DIR=$(OUTPUT_DIR) NODE_OPTIONS=--experimental-vm-modules NODE_NO_WARNINGS=1 \
         $(shell npm bin)/jest --detectOpenHandles -i
 
 $(OUTPUT_DIR)/wasm-table.ts: $(WASI_TEST_SUITE_DST_TESTS) $(WASMTIME_DST_TESTS)
@@ -51,7 +61,7 @@ $(MEMFS_DST): $(MEMFS_SRC)
 
 $(BUNDLE): $(wildcard ../dist/**) $(wildcard ./driver/**) $(MEMFS_DST) $(OUTPUT_DIR)/wasm-table.ts node_modules
 	mkdir -p $(@D)
-	$(shell npm bin)/esbuild --bundle ./driver/index.ts --outfile=$@ --format=esm --log-level=warning --external:*.wasm
+	$(shell npm bin)/esbuild --bundle ./driver/worker.ts --outfile=$@ --format=esm --log-level=warning --external:*.wasm
 
 $(WASM_OPT):
 	@$(call color,"downloading binaryen")
@@ -60,3 +70,12 @@ $(WASM_OPT):
 	echo '9057c8f3f0bbfec47a95985c8f0faad8cc2aa3932e94a7d6b705e245ed140e19  binaryen.tar.gz' | sha256sum -c
 	tar zxvf binaryen.tar.gz --strip-components=1 --touch -C ../build/binaryen
 	rm binaryen.tar.gz
+
+export WASI_CC      := $(abspath ${WASI_SDK_PATH}/bin/clang) -target wasm32-wasi --sysroot=${WASI_SYSROOT}
+export WASI_CFLAGS  := -Oz -flto
+export WASI_LDFLAGS := -flto -Wl,--allow-undefined
+
+$(BENCHMARK_DST)/%.wasm: $(WASI_SDK_PATH) $(WASM_OPT) $(BENCHMARK_SRC)/%.c
+	mkdir -p $(BENCHMARK_DST)
+	$(WASI_CC) $(WASI_CFLAGS) $(WASI_LDFLAGS) subjects/$*.c -o $(BENCHMARK_DST)/$*.wasm 
+	$(WASM_OPT) -g -O --asyncify $(BENCHMARK_DST)/$*.wasm -o $(BENCHMARK_DST)/$*.asyncify.wasm

test/benchmark-build.mjs

@@ -0,0 +1,43 @@
+// @ts-check
+import * as esbuild from 'esbuild'
+import * as path from 'path'
+
+const OUT_DIR = '../build/test/'
+
+/**
+ * Node currently can't import wasm files in the way that we do with Workers, it either throws an
+ * error if you try to import a wasm file or imports an instantiated wasm instance. Whereas in
+ * Workers we get a WebAssembly.Module as the default export if we import a wasm file, so this
+ * plugin is to replicate that behavior in the bundle.
+ * @type {import('esbuild').Plugin}
+ */
+const wasmLoaderPlugin = {
+  name: 'wasm-module-loader',
+  setup: (build) => {
+    build.onResolve({ filter: /\.wasm$/ }, (args) => ({
+      path: args.path,
+      namespace: 'wasm-module',
+    }))
+
+    build.onLoad({ filter: /.*/, namespace: 'wasm-module' }, (args) => ({
+      contents: `
+        import * as fs from 'fs';
+        import * as path from 'path';
+        import * as url from 'url';
+
+        const __dirname = path.dirname(url.fileURLToPath(import.meta.url));
+        export default new WebAssembly.Module(fs.readFileSync(path.resolve(__dirname, '${args.path}')));
+      `,
+    }))
+  },
+}
+
+esbuild.build({
+  bundle: true,
+  outfile: path.join(OUT_DIR, 'standalone.mjs'),
+  format: 'esm',
+  logLevel: 'warning',
+  entryPoints: ['./driver/standalone.ts'],
+  plugins: [wasmLoaderPlugin],
+  platform: 'node',
+})

test/benchmark.test.ts

@@ -0,0 +1,53 @@
+import * as child from 'node:child_process'
+import * as fs from 'node:fs'
+import { cwd } from 'node:process'
+import path from 'path/posix'
+import type { ExecOptions } from './driver/common'
+
+const { OUTPUT_DIR } = process.env
+const moduleNames = fs
+  .readdirSync(`${OUTPUT_DIR}/benchmark`)
+  .map((dirent) => `benchmark/${dirent}`)
+
+for (const modulePath of moduleNames) {
+  const prettyName = modulePath.split('/').pop()
+  if (!prettyName) throw new Error('unreachable')
+
+  test(`${prettyName}`, async () => {
+    const execOptions: ExecOptions = {
+      moduleName: prettyName,
+      asyncify: prettyName.endsWith('.asyncify.wasm'),
+      fs: {},
+      preopens: [],
+      returnOnExit: false,
+    }
+
+    // Spawns a child process that runs the wasm so we can isolate the profiling to just that
+    // specific test case.
+    const proc = child.execFile(
+      `node`,
+      [
+        '--experimental-vm-modules',
+        '--cpu-prof',
+        '--cpu-prof-dir=./prof',
+        `--cpu-prof-name=${prettyName}.${Date.now()}.cpuprofile`,
+        'standalone.mjs',
+        modulePath,
+        JSON.stringify(execOptions),
+      ],
+      {
+        encoding: 'utf8',
+        cwd: OUTPUT_DIR,
+      }
+    )
+
+    let stderr = ''
+    proc.stderr?.on('data', (data) => (stderr += data))
+
+    const exitCode = await new Promise((resolve) => proc.once('exit', resolve))
+
+    if (exitCode !== 0) {
+      console.error(`Child process exited with code ${exitCode}:\n${stderr}`)
+    }
+  })
+}

test/driver/common.ts

@@ -1,4 +1,4 @@
-import { Environment, _FS } from '@cloudflare/workers-wasi'
+import { Environment, WASI, _FS } from '@cloudflare/workers-wasi'
 
 export interface ExecOptions {
   args?: string[]
@@ -13,5 +13,78 @@ export interface ExecOptions {
 
 export interface ExecResult {
   stdout: string
+  stderr: string
   status?: number
 }
+
+export const exec = async (
+  options: ExecOptions,
+  wasm: WebAssembly.Module,
+  body?: ReadableStream<Uint8Array>
+): Promise<ExecResult> => {
+  let TransformStream = global.TransformStream
+
+  if (TransformStream === undefined) {
+    const streams = await import('node:stream/web').catch(() => {
+      throw new Error('unreachable')
+    })
+
+    TransformStream = streams.TransformStream as typeof TransformStream
+  }
+
+  const stdout = new TransformStream()
+  const stderr = new TransformStream()
+
+  const wasi = new WASI({
+    args: options.args,
+    env: options.env,
+    fs: options.fs,
+    preopens: options.preopens,
+    returnOnExit: options.returnOnExit,
+    stderr: stderr.writable,
+    stdin: body,
+    stdout: stdout.writable,
+    streamStdio: options.asyncify,
+  })
+  const instance = new WebAssembly.Instance(wasm, {
+    wasi_snapshot_preview1: wasi.wasiImport,
+  })
+  const promise = wasi.start(instance)
+
+  const streams = await Promise.all([
+    collectStream(stdout.readable),
+    collectStream(stderr.readable),
+  ])
+
+  try {
+    const result = {
+      stdout: streams[0],
+      stderr: streams[1],
+      status: await promise,
+    }
+    return result
+  } catch (e: any) {
+    e.message = `${e}\n\nstdout:\n${streams[0]}\n\nstderr:\n${streams[1]}\n\n`
+    throw e
+  }
+}
+
+const collectStream = async (stream: ReadableStream): Promise<string> => {
+  const chunks: Uint8Array[] = []
+
+  // @ts-ignore
+  for await (const chunk of stream) {
+    chunks.push(chunk)
+  }
+
+  const size = chunks.reduce((acc, chunk) => acc + chunk.byteLength, 0)
+  const buffer = new Uint8Array(size)
+  let offset = 0
+
+  chunks.forEach((chunk) => {
+    buffer.set(chunk, offset)
+    offset += chunk.byteLength
+  })
+
+  return new TextDecoder().decode(buffer)
+}

test/driver/standalone.ts

@@ -0,0 +1,29 @@
+import * as fs from 'node:fs/promises'
+import { ReadableStream } from 'node:stream/web'
+import { exec } from './common'
+
+const [modulePath, rawOptions] = process.argv.slice(2)
+const options = JSON.parse(rawOptions)
+
+const nulls = new Uint8Array(4096).fill(0)
+let written = 0
+
+const stdinStream = new ReadableStream<Uint8Array>({
+  pull: (controller) => {
+    if (written > 1_000_000) {
+      controller.close()
+    } else {
+      controller.enqueue(nulls)
+      written += nulls.byteLength
+    }
+  },
+})
+
+fs.readFile(modulePath)
+  .then((wasmBytes) => new WebAssembly.Module(wasmBytes))
+  .then((wasmModule) => exec(options, wasmModule, stdinStream as any))
+  .then((result) => {
+    console.log(result.stdout)
+    console.error(result.stderr)
+    process.exit(result.status)
+  })

test/driver/worker.ts

@@ -0,0 +1,17 @@
+import { ModuleTable } from '../../build/test/wasm-table'
+import { ExecOptions, exec } from './common'
+
+export default {
+  async fetch(request: Request) {
+    const options: ExecOptions = JSON.parse(
+      atob(request.headers.get('EXEC_OPTIONS')!)
+    )
+
+    const result = await exec(
+      options,
+      ModuleTable[options.moduleName],
+      request.body ?? undefined
+    )
+    return new Response(JSON.stringify(result))
+  },
+}

test/driver/wrangler.toml

@@ -17,4 +17,3 @@ dir = "../../build/test/"
 [[build.upload.rules]]
 type = "CompiledWasm"
 globs = ["**/*.wasm"]
-

test/generate-wasm-table.mjs

@@ -1,7 +1,5 @@
-// generate a table with explicit imports since we can't dynamically load
-// WebAssembly modules in our worker
-import path from 'path'
-import fs from 'fs/promises'
+import * as path from 'path'
+import * as fs from 'fs/promises'
 
 const recurseFiles = async (dir) => {
   const entries = await fs.readdir(dir, { withFileTypes: true })
@@ -21,19 +19,24 @@ const recurseFiles = async (dir) => {
   ).flat()
 }
 
-
 const dir = path.resolve(process.argv[2])
-const files = (await recurseFiles(dir)).map(f => path.join('./', f.substr(dir.length)))
+const files = (await recurseFiles(dir)).map((f) =>
+  path.join('./', f.substr(dir.length))
+)
+
 for (const file of files) {
   console.log('// @ts-ignore')
-  const identifier = file.replace(/[-\/\.]/g, "_")
+  const identifier = file.replace(/[-\/\.]/g, '_')
   console.log(`import ${identifier} from '${file}'`)
 }
 console.log()
 
-console.log('export const ModuleTable: { [key: string]: WebAssembly.Module } = {')
+console.log(
+  'export const ModuleTable: { [key: string]: WebAssembly.Module } = {'
+)
 for (const file of files) {
-  const identifier = file.replace(/[-\/\.]/g, "_")
+  const identifier = file.replace(/[-\/\.]/g, '_')
   console.log(`  '${file}': ${identifier},`)
 }
+
 console.log('}')