modified run-wasm.mjs that attempt to mitigate the blocking read of stdin 🫠

run-wasm-worker.mjs

#!/usr/bin/env -S node --disable-warning=ExperimentalWarning --max-old-space-size=65536 --no-turbo-fast-api-calls --wasm-lazy-validation

import fs from "node:fs/promises";
import { WASI } from "node:wasi";

function parseArgv(args) {
  const i = args.indexOf("-0");
  return i === -1 ? args : args.slice(i + 2);
}

const argv = parseArgv(process.argv.slice(2));

let stdinIsAavailable = false
let stdinGlobalBuffer = []
process.stdin.on('readable', () => {
  console.warn('READABLE')
  stdinIsAavailable = true
})

/** @type {WebAssembly.Instance} */
let ginstance

function getMemory() {
  return /** @type {WebAssembly.Memory} */(ginstance.exports.memory).buffer
}

function readMemoryInt32(addr) {
  const memory = getMemory()
  const dv = new DataView(memory)
  return dv.getInt32(addr, true)
}

function writeMemoryUint32(addr, n) {
  const memory = getMemory()
  const dv = new DataView(memory)
  return dv.setUint32(addr, n, true)
}

function iovsTotalLen(iovs, cnt) {
  const memory = getMemory()
  const dv = new DataView(memory)
  let res = 0
  for (let i = 0; i < cnt; i++) {
    const len = dv.getUint32(iovs + i * 8 + 4, true)
    res += len
  }
  return res
}

function collectIovs(iovs, cnt, cap) {
  const memory = getMemory()
  const uarr = new Uint8Array(memory)
  const dv = new DataView(memory)
  let res = []
  for (let i = 0; i < cnt; i++) {
    const ptr = dv.getUint32(iovs + i * 8, true)
    const len = dv.getUint32(iovs + i * 8 + 4, true)
    const buf = Array.from(uarr.slice(ptr, ptr + len))
    res = res.concat(buf)
  }
  const rbuf = new Uint8Array(res)
  return new TextDecoder().decode(cap != null ? rbuf.slice(0, cap) : rbuf)
}

/** @param {Uint8Array} data */
function fillIovs(iovs, cnt, data) {
  const memory = getMemory()
  const uarr = new Uint8Array(memory)
  const dv = new DataView(memory)

  let offset = 0

  for (let i = 0; i < cnt && offset < data.length; i++) {
    const ptr = dv.getUint32(iovs + i * 8, true)
    const len = dv.getUint32(iovs + i * 8 + 4, true)

    const chunkSize = Math.min(len, data.length - offset)
    uarr.set(data.subarray(offset, offset + chunkSize), ptr)
    offset += chunkSize
  }

  return data.length - offset
}

function patchImports(imports) {
  const overrides = {
    fd_write(_k, fn) {
      return (fd, iovs, cnt, res) => {
        console.warn(`[WASI WRITE] fd=${fd} src=${JSON.stringify(collectIovs(iovs, cnt))}`)
        const ret = fn(fd, iovs, cnt, res)
        if (ret == 0) {
          console.warn(`  --> DONE len=${readMemoryInt32(res)}`)
        } else {
          console.warn(`  --> FAIL ret=${ret}`)
        }
        return ret
      }
    },
    fd_read(_k, fn) {
      return (fd, iovs, cnt, res) => {
        console.warn(`\x1b[1;33m[WASI READ]\x1b[0m fd=${fd} nread=${iovsTotalLen(iovs, cnt)}`)
        if (fd == 0) {
          // console.warn('buffers', stdinBuffers)
          // if (!stdinBuffers.length) {
          if (!stdinIsAavailable) {
            console.warn('  --> FAIL, buffer is empty')
            return 6  /* EAGAIN */
          }
          const want = iovsTotalLen(iovs, cnt)

          while (stdinGlobalBuffer.length < want) {
            const chunk = process.stdin.read()
            if (chunk == null) {
              stdinIsAavailable = false
              break
            }
            stdinGlobalBuffer = stdinGlobalBuffer.concat(Array.from(chunk))
          }

          const remaining = fillIovs(iovs, cnt, new Uint8Array(stdinGlobalBuffer))

          const filled = stdinGlobalBuffer.length - remaining

          stdinGlobalBuffer = stdinGlobalBuffer.slice(filled)
          writeMemoryUint32(res, filled)

          return 0
        }
        const ret = fn(fd, iovs, cnt, res)
        if (ret == 0) {
          console.warn(`  --> DONE len=${readMemoryInt32(res)} dest=${JSON.stringify(collectIovs(iovs, cnt, readMemoryInt32(res)))}`)
        } else {
          console.warn(`  --> FAIL ret=${ret}`)
        }
        return ret
      }
    },
    poll_oneoff(_k, fn) {
      return (sin, sout, nsub, nevtptr) => {
        console.warn(`[WASI poll_oneoff] nsub=${nsub}`)
        const subs = []
        for (let i = 0; i < nsub; i++) {
          const memory = getMemory()
          const dv = new DataView(memory)

          const userdata = dv.getBigUint64(sin + i * 48, true)
          const tp = dv.getInt8(sin + i * 48 + 8, true)
          if (tp == 0  /* clock */) {
            const clockid = dv.getUint32(sin + i * 48 + 16, true)
            const timeout = dv.getBigUint64(sin + i * 48 + 24, true)
            const subclockflags = dv.getUint16(sin + i * 48 + 40, true)
            subs.push({ type: 'clock', clockid, timeout, subclockflags, userdata })
          } else if (tp == 1  /* read */ || tp == 2) {
            const tpIsRead = tp == 1
            const fd = dv.getUint32(sin + i * 48 + 16, true)
            subs.push({ type: tpIsRead ? 'read' : 'write', fd, userdata })
          }
        }

        console.warn('SUBS=', subs)
        if (subs.length == 1 && subs[0].type == 'read' && subs[0].fd == 0) {
          if (stdinIsAavailable) {
            // CIRCUMVENT the blocking stdin polling
            dv.setUint32(sout, subs[0].userdata, true)
            dv.setUint16(sout + 8, 0, true)  // errno
            dv.setUint16(sout + 10, 1, true)  // tp
            dv.setBigInt64(sout + 16, BigInt(stdinGlobalBuffer.length), true)  // size
            dv.setUint16(sout + 24, 0, true)  // hangup?

            writeMemoryUint32(nevtptr, 1)
            return 0
          }
          // not knowing what to do :(
        }

        const ret = fn(sin, sout, nsub, nevtptr)
        console.warn(`  RET=${ret} nevt=${readMemoryInt32(nevtptr)}`)
        return ret
      }
    },
  }

  function fallback(k, fn) {
    return (...args) => {
      console.warn(`[WASI ${k}] ${args}`)
      const ret = fn(...args)
      console.warn(`  RET=${ret}`)
      return ret
    }
  }

  Object.keys(imports).forEach(k => {
    imports[k] = (overrides[k] ?? fallback)(k, imports[k])
  })
  return imports
}

async function main() {
  const wasi = new WASI({
    version: "preview1",
    args: ['als.wasm', ...argv],
    env: { PATH: "", PWD: process.cwd() },
    preopens: { "/": "/" },
  });

  const instance = (
    await WebAssembly.instantiate(await fs.readFile('./als.wasm'), {
      wasi_snapshot_preview1: patchImports(wasi.wasiImport),
    })
  ).instance;

  ginstance = instance

  try {
    let ec = wasi.start(instance);
    // wasmtime reports "exit with invalid exit status outside of [0..126)"
    if (ec >= 126) {
      ec = 1;
    }
    process.exit(ec);
  } catch (err) {
    if (!(err instanceof WebAssembly.RuntimeError)) {
      throw err;
    }
    console.error(err.stack);
    // wasmtime exits with 128+SIGABRT
    process.exit(134);
  }
}

main()