perf: rm fixed buffer in favor of alloc

This commit is contained in:
Nurul Huda (Apon) 2026-07-18 13:44:46 +06:00
parent 8773c00cd2
commit c1d06113b5
No known key found for this signature in database
GPG key ID: 5D3F1DE2855A2F79
10 changed files with 151 additions and 88 deletions

View file

@ -2,6 +2,7 @@ import { ZxBridge } from "../wasm";
import { createKVImports, createMemoryKV } from "../kv";
import { createFetchImports } from "../fetch";
import { createD1Imports } from "../db";
import { bindWasmAlloc, type WasmAllocRef } from "../wasm/core";
import { createWasiImports } from "../wasi";
import { buildWsImports, attachWebSocket } from "../runtime";
import type { WsState } from "../runtime";
@ -122,13 +123,14 @@ export function createWebSocketDO(
const kvBindings = options?.kv?.(this.env);
const dbBindings = options?.db?.(this.env);
const allocRef: WasmAllocRef = { current: null };
const instance = new WebAssembly.Instance(module, {
wasi_snapshot_preview1: wasiImport,
__zx_sys: sysImports,
__zx_ws: wsImports,
__zx_kv: createKVImports(kvBindings ?? { default: createMemoryKV() }, mem),
__zx_db: createD1Imports(dbBindings ?? {}, mem),
__zx_kv: createKVImports(kvBindings ?? { default: createMemoryKV() }, mem, allocRef),
__zx_db: createD1Imports(dbBindings ?? {}, mem, allocRef),
__zx_net: createFetchImports(mem),
...(options?.imports ? options.imports(mem) : {}),
...bridgeImports,
@ -137,6 +139,7 @@ export function createWebSocketDO(
wasmMemory = instance.exports.memory as WebAssembly.Memory;
setMemory(wasmMemory);
bridgeRef.current = new ZxBridge(instance.exports);
bindWasmAlloc(allocRef, instance.exports);
const start = (WebAssembly as any).promising(instance.exports._start as Function);

View file

@ -1,3 +1,7 @@
import { type WasmAllocRef, writeJsonOut } from "./wasm/core";
export type { WasmAllocRef };
export type D1Value =
| null
| string
@ -118,21 +122,14 @@ function valuesToWireRows(rows: unknown[][]): WireValue[][] {
export function createD1Imports(
bindings: Record<string, D1Database>,
getMemory: () => WebAssembly.Memory,
allocRef: WasmAllocRef,
): Record<string, unknown> {
const encoder = new TextEncoder();
const decoder = new TextDecoder();
function readStr(ptr: number, len: number): string {
return decoder.decode(new Uint8Array(getMemory().buffer, ptr, len));
}
function writeJson(buf_ptr: number, buf_max: number, value: unknown): number {
const data = encoder.encode(JSON.stringify(value));
if (data.length > buf_max) return -2;
new Uint8Array(getMemory().buffer, buf_ptr, data.length).set(data);
return data.length;
}
function binding(ns: string): D1Database | null {
return bindings[ns] ?? bindings["default"] ?? null;
}
@ -157,10 +154,10 @@ export function createD1Imports(
if (typeof Suspending !== "function") {
return {
db_open: (_ns: number, _ns_len: number): number => -1,
db_run: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number, _h: number): number => -1,
db_get: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number, _h: number): number => -1,
db_all: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number, _h: number): number => -1,
db_values: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number, _h: number): number => -1,
db_run: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number): number => -1,
db_get: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number): number => -1,
db_all: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number): number => -1,
db_values: (_a: number, _b: number, _c: number, _d: number, _e: number, _f: number, _g: number): number => -1,
};
}
@ -171,12 +168,12 @@ export function createD1Imports(
ns_ptr: number, ns_len: number,
sql_ptr: number, sql_len: number,
bindings_ptr: number, bindings_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const stmt = await statement(ns_ptr, ns_len, sql_ptr, sql_len, bindings_ptr, bindings_len);
if (!stmt) return -1;
const result = await stmt.run();
return writeJson(buf_ptr, buf_max, {
return writeJsonOut(getMemory, allocRef, out_ptr, {
last_insert_rowid: result.meta?.last_row_id ?? 0,
changes: result.meta?.changes ?? 0,
});
@ -186,37 +183,49 @@ export function createD1Imports(
ns_ptr: number, ns_len: number,
sql_ptr: number, sql_len: number,
bindings_ptr: number, bindings_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const stmt = await statement(ns_ptr, ns_len, sql_ptr, sql_len, bindings_ptr, bindings_len);
if (!stmt) return -1;
const row = await stmt.first<Record<string, unknown>>();
if (!row) return 0;
return writeJson(buf_ptr, buf_max, [objectToWireRow(row)]);
if (!row) {
new DataView(getMemory().buffer).setUint32(out_ptr, 0, true);
return 0;
}
return writeJsonOut(getMemory, allocRef, out_ptr, [objectToWireRow(row)]);
}),
db_all: new Suspending(async (
ns_ptr: number, ns_len: number,
sql_ptr: number, sql_len: number,
bindings_ptr: number, bindings_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const stmt = await statement(ns_ptr, ns_len, sql_ptr, sql_len, bindings_ptr, bindings_len);
if (!stmt) return -1;
const result = await stmt.all<Record<string, unknown>>();
return writeJson(buf_ptr, buf_max, (result.results ?? []).map(objectToWireRow));
const rows = (result.results ?? []).map(objectToWireRow);
if (rows.length === 0) {
new DataView(getMemory().buffer).setUint32(out_ptr, 0, true);
return 0;
}
return writeJsonOut(getMemory, allocRef, out_ptr, rows);
}),
db_values: new Suspending(async (
ns_ptr: number, ns_len: number,
sql_ptr: number, sql_len: number,
bindings_ptr: number, bindings_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const stmt = await statement(ns_ptr, ns_len, sql_ptr, sql_len, bindings_ptr, bindings_len);
if (!stmt) return -1;
const rows = await stmt.raw();
return writeJson(buf_ptr, buf_max, valuesToWireRows(rows));
if (rows.length === 0) {
new DataView(getMemory().buffer).setUint32(out_ptr, 0, true);
return 0;
}
return writeJsonOut(getMemory, allocRef, out_ptr, valuesToWireRows(rows));
}),
};
}

View file

@ -1,3 +1,7 @@
import { type WasmAllocRef, writeBytesOut, writeJsonOut } from "./wasm/core";
export type { WasmAllocRef };
// Minimal type definition for a key-value namespace
export interface KVNamespace {
get(key: string): Promise<string | null>;
@ -82,6 +86,7 @@ function isSyncKVNamespace(binding: KVNamespace): binding is SyncKVNamespace {
export function createKVImports(
bindings: Record<string, KVNamespace>,
getMemory: () => WebAssembly.Memory,
allocRef: WasmAllocRef,
): Record<string, unknown> {
const encoder = new TextEncoder();
const decoder = new TextDecoder();
@ -90,12 +95,6 @@ export function createKVImports(
return decoder.decode(new Uint8Array(getMemory().buffer, ptr, len));
}
function writeBytes(buf_ptr: number, buf_max: number, data: Uint8Array): number {
if (data.length > buf_max) return -2;
new Uint8Array(getMemory().buffer, buf_ptr, data.length).set(data);
return data.length;
}
function binding(ns: string): KVNamespace | null {
return bindings[ns] ?? bindings["default"] ?? null;
}
@ -108,12 +107,12 @@ export function createKVImports(
}
return {
kv_get: (ns_ptr: number, ns_len: number, key_ptr: number, key_len: number, buf_ptr: number, buf_max: number): number => {
kv_get: (ns_ptr: number, ns_len: number, key_ptr: number, key_len: number, out_ptr: number): number => {
const b = syncBinding(readStr(ns_ptr, ns_len));
if (!b) return -1;
const value = b.getSync(readStr(key_ptr, key_len));
if (value === null) return -1;
return writeBytes(buf_ptr, buf_max, encoder.encode(value));
return writeBytesOut(getMemory, allocRef, out_ptr, encoder.encode(value));
},
kv_put: (ns_ptr: number, ns_len: number, key_ptr: number, key_len: number, val_ptr: number, val_len: number, ttl_seconds: number): number => {
const b = syncBinding(readStr(ns_ptr, ns_len));
@ -127,12 +126,12 @@ export function createKVImports(
b.deleteSync(readStr(key_ptr, key_len));
return 0;
},
kv_list: (ns_ptr: number, ns_len: number, pfx_ptr: number, pfx_len: number, buf_ptr: number, buf_max: number): number => {
kv_list: (ns_ptr: number, ns_len: number, pfx_ptr: number, pfx_len: number, out_ptr: number): number => {
const b = syncBinding(readStr(ns_ptr, ns_len));
if (!b) return writeBytes(buf_ptr, buf_max, encoder.encode("[]"));
if (!b) return writeJsonOut(getMemory, allocRef, out_ptr, []);
const prefix = readStr(pfx_ptr, pfx_len);
const result = b.listSync(prefix.length > 0 ? { prefix } : undefined);
return writeBytes(buf_ptr, buf_max, encoder.encode(JSON.stringify(result.keys.map((k) => k.name))));
return writeJsonOut(getMemory, allocRef, out_ptr, result.keys.map((k) => k.name));
},
};
}
@ -141,13 +140,13 @@ export function createKVImports(
kv_get: new Suspending(async (
ns_ptr: number, ns_len: number,
key_ptr: number, key_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const b = binding(readStr(ns_ptr, ns_len));
if (!b) return -1;
const value = await b.get(readStr(key_ptr, key_len));
if (value === null) return -1;
return writeBytes(buf_ptr, buf_max, encoder.encode(value));
return writeBytesOut(getMemory, allocRef, out_ptr, encoder.encode(value));
}),
kv_put: new Suspending(async (
@ -175,13 +174,18 @@ export function createKVImports(
kv_list: new Suspending(async (
ns_ptr: number, ns_len: number,
prefix_ptr: number, prefix_len: number,
buf_ptr: number, buf_max: number,
out_ptr: number,
): Promise<number> => {
const b = binding(readStr(ns_ptr, ns_len));
if (!b) return writeBytes(buf_ptr, buf_max, encoder.encode("[]"));
if (!b) return writeJsonOut(getMemory, allocRef, out_ptr, []);
const prefix = readStr(prefix_ptr, prefix_len);
const result = await b.list(prefix.length > 0 ? { prefix } : undefined);
return writeBytes(buf_ptr, buf_max, encoder.encode(JSON.stringify(result.keys.map((k) => k.name))));
const names = result.keys.map((k) => k.name);
if (names.length === 0) {
new DataView(getMemory().buffer).setUint32(out_ptr, 0, true);
return 0;
}
return writeJsonOut(getMemory, allocRef, out_ptr, names);
}),
};
}

View file

@ -2,6 +2,7 @@ import { ZxWasiBridge } from "./wasm/wasi";
import { createKVImports, createMemoryKV } from "./kv";
import { createFetchImports } from "./fetch";
import { createD1Imports } from "./db";
import { bindWasmAlloc, type WasmAllocRef } from "./wasm/core";
import { createWasiImports, ProcExit, mergeUint8Arrays } from "./wasi";
import type { WASI } from "./wasi";
import type { KVNamespace } from "./kv";
@ -263,6 +264,7 @@ export async function run({
const mem = () => wasmMemory;
const bridgeRef: { current: ZxWasiBridge | null } = { current: null };
const allocRef: WasmAllocRef = { current: null };
const Suspending = (WebAssembly as any).Suspending;
const jspi = typeof Suspending === 'function';
@ -279,8 +281,8 @@ export async function run({
wasi_snapshot_preview1: { ...wasi?.wasiImport, ...wasiImport },
__zx_sys: buildSysImports(jspi, Suspending),
__zx_ws: buildWsImports(jspi ? Suspending : null, mem, new TextDecoder(), wsState),
__zx_kv: createKVImports(kvBindings ?? { default: createMemoryKV() }, mem),
__zx_db: createD1Imports(dbBindings ?? {}, mem),
__zx_kv: createKVImports(kvBindings ?? { default: createMemoryKV() }, mem, allocRef),
__zx_db: createD1Imports(dbBindings ?? {}, mem, allocRef),
__zx_net: createFetchImports(mem),
...(imports ? imports(mem) : {}),
...ZxWasiBridge.createImportObject(bridgeRef),
@ -289,6 +291,7 @@ export async function run({
wasmMemory = instance.exports.memory as WebAssembly.Memory;
setMemory(wasmMemory);
bridgeRef.current = new ZxWasiBridge(instance.exports);
bindWasmAlloc(allocRef, instance.exports);
const wasmPromise = executeWasm(instance, jspi, Suspending, wsState);

View file

@ -322,3 +322,37 @@ export class ZxBridgeCore {
};
}
}
export type WasmAllocRef = { current: ((size: number) => number) | null };
export function writeBytesOut(
getMemory: () => WebAssembly.Memory,
allocRef: WasmAllocRef,
outPtrAddr: number,
data: Uint8Array,
): number {
if (data.length === 0) {
new DataView(getMemory().buffer).setUint32(outPtrAddr, 0, true);
return 0;
}
const alloc = allocRef.current;
if (!alloc) return -1;
const ptr = alloc(data.length);
if (!ptr) return -1;
const buffer = getMemory().buffer;
new Uint8Array(buffer, ptr, data.length).set(data);
new DataView(buffer).setUint32(outPtrAddr, ptr, true);
return data.length;
}
export function writeJsonOut(
getMemory: () => WebAssembly.Memory,
allocRef: WasmAllocRef,
outPtrAddr: number,
value: unknown,
): number {
return writeBytesOut(getMemory, allocRef, outPtrAddr, textEncoder.encode(JSON.stringify(value)));
}
export function bindWasmAlloc(allocRef: WasmAllocRef, exports: WebAssembly.Exports): void {
allocRef.current = exports.__zx_alloc as (size: number) => number;
}

View file

@ -8,9 +8,12 @@ export {
getMemoryView,
readString,
writeBytes,
writeBytesOut,
writeJsonOut,
bindWasmAlloc,
ZxBridgeCore,
} from "./core";
export type { CallbackTypeValue } from "./core";
export type { CallbackTypeValue, WasmAllocRef } from "./core";
import {
ZxBridgeCore,
@ -26,6 +29,7 @@ import {
getMemoryView,
} from "./core";
import { createKVImports, type KVNamespace } from "../kv";
import { bindWasmAlloc, type WasmAllocRef } from "./core";
import { createFetchImports } from "../fetch";
import { createBrowserKVBindings } from "../browser/kv";
import type {
@ -735,12 +739,13 @@ export async function init(options: InitOptions = {}): Promise<{ source: WebAsse
let wasmMemory: WebAssembly.Memory | null = null;
const kvBindings = options.kv ?? createBrowserKVBindings();
const allocRef: WasmAllocRef = { current: null };
const kvImportObject = {
__zx_kv: createKVImports(kvBindings, () => {
if (wasmMemory) return wasmMemory;
if (jsz.memory) return jsz.memory;
throw new Error("WASM memory is not ready");
}),
}, allocRef),
};
const importObject = Object.assign(
@ -755,6 +760,7 @@ export async function init(options: InitOptions = {}): Promise<{ source: WebAsse
wasmMemory = instance.exports.memory as WebAssembly.Memory;
jsz.memory = wasmMemory;
bindWasmAlloc(allocRef, instance.exports);
const bridge = new ZxBridge(instance.exports);
bridgeRef.current = bridge;

View file

@ -424,6 +424,12 @@ pub fn initInner(
zx_module.addImport("app", app_module);
exe.root_module.addImport("app", app_module);
exe.root_module.addImport("zx", zx_module);
if (exe.root_module.resolved_target) |t| {
if (t.result.os.tag == .wasi) {
exe.rdynamic = true;
exe.export_memory = true;
}
}
exe.step.dependOn(&transpile_cmd.step);
exe.step.name = b.fmt("install server exe", .{});

View file

@ -393,3 +393,14 @@ pub fn logFn(
defer std.heap.wasm_allocator.free(msg);
ext._log(level, msg.ptr, msg.len);
}
export fn __zx_alloc(size: usize) ?[*]u8 {
if (size == 0) return null;
const ptr = std.heap.wasm_allocator.alloc(u8, size) catch return null;
return ptr.ptr;
}
export fn __zx_free(ptr: [*]u8, size: usize) void {
if (size == 0) return;
std.heap.wasm_allocator.free(ptr[0..size]);
}

View file

@ -180,7 +180,7 @@ fn runQuery(binding_name: []const u8, sql: []const u8, bindings: Db.Bindings) !D
defer bindings_writer.deinit();
try writeBindingsJson(&bindings_writer.writer, bindings);
var buf: [8192]u8 = undefined;
var ptr: [*]u8 = undefined;
const n = ext.db_run(
binding_name.ptr,
binding_name.len,
@ -188,12 +188,13 @@ fn runQuery(binding_name: []const u8, sql: []const u8, bindings: Db.Bindings) !D
sql.len,
bindings_writer.written().ptr,
bindings_writer.written().len,
&buf,
buf.len,
&ptr,
);
if (n < 0) return error.DatabaseRunFailed;
if (n == 0) return error.DatabaseRunFailed;
defer std.heap.wasm_allocator.free(ptr[0..@intCast(n)]);
const parsed = try std.json.parseFromSlice(WireRunResult, std.heap.wasm_allocator, buf[0..@intCast(n)], .{
const parsed = try std.json.parseFromSlice(WireRunResult, std.heap.wasm_allocator, ptr[0..@intCast(n)], .{
.ignore_unknown_fields = true,
.allocate = .alloc_always,
});
@ -216,7 +217,7 @@ fn selectRows(
defer bindings_writer.deinit();
try writeBindingsJson(&bindings_writer.writer, bindings);
var buf: [65536]u8 = undefined;
var ptr: [*]u8 = undefined;
const n = op(
binding_name.ptr,
binding_name.len,
@ -224,13 +225,13 @@ fn selectRows(
sql.len,
bindings_writer.written().ptr,
bindings_writer.written().len,
&buf,
buf.len,
&ptr,
);
if (n < 0) return error.DatabaseQueryFailed;
if (n == 0) return &[_]Db.Row{};
defer std.heap.wasm_allocator.free(ptr[0..@intCast(n)]);
const parsed = try std.json.parseFromSlice([]WireRow, allocator, buf[0..@intCast(n)], .{
const parsed = try std.json.parseFromSlice([]WireRow, allocator, ptr[0..@intCast(n)], .{
.ignore_unknown_fields = true,
.allocate = .alloc_always,
});
@ -244,7 +245,7 @@ fn selectValues(allocator: std.mem.Allocator, binding_name: []const u8, sql: []c
defer bindings_writer.deinit();
try writeBindingsJson(&bindings_writer.writer, bindings);
var buf: [65536]u8 = undefined;
var ptr: [*]u8 = undefined;
const n = ext.db_values(
binding_name.ptr,
binding_name.len,
@ -252,13 +253,13 @@ fn selectValues(allocator: std.mem.Allocator, binding_name: []const u8, sql: []c
sql.len,
bindings_writer.written().ptr,
bindings_writer.written().len,
&buf,
buf.len,
&ptr,
);
if (n < 0) return error.DatabaseQueryFailed;
if (n == 0) return &[_][]const Db.Value{};
defer std.heap.wasm_allocator.free(ptr[0..@intCast(n)]);
const parsed = try std.json.parseFromSlice([][]WireValue, allocator, buf[0..@intCast(n)], .{
const parsed = try std.json.parseFromSlice([][]WireValue, allocator, ptr[0..@intCast(n)], .{
.ignore_unknown_fields = true,
.allocate = .alloc_always,
});
@ -370,8 +371,7 @@ const ext = struct {
sql_len: usize,
bindings_ptr: [*]const u8,
bindings_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
pub extern "__zx_db" fn db_get(
@ -381,8 +381,7 @@ const ext = struct {
sql_len: usize,
bindings_ptr: [*]const u8,
bindings_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
pub extern "__zx_db" fn db_all(
@ -392,8 +391,7 @@ const ext = struct {
sql_len: usize,
bindings_ptr: [*]const u8,
bindings_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
pub extern "__zx_db" fn db_values(
@ -403,8 +401,7 @@ const ext = struct {
sql_len: usize,
bindings_ptr: [*]const u8,
bindings_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
};

View file

@ -3,24 +3,13 @@ const Wasm = @This();
const std = @import("std");
const Kv = @import("../Kv.zig");
const max_response_bytes = 16 * 1024 * 1024;
fn get(_: ?*anyopaque, ns: []const u8, allocator: std.mem.Allocator, key: []const u8) !?[]u8 {
var capacity: usize = 65536;
var buf = try allocator.alloc(u8, capacity);
defer allocator.free(buf);
while (true) {
const n = ext.kv_get(ns.ptr, ns.len, key.ptr, key.len, buf.ptr, capacity);
if (n == -2) {
capacity *= 2;
if (capacity > max_response_bytes) return error.InvalidResponse;
buf = try allocator.realloc(buf, capacity);
continue;
}
if (n < 0) return null;
return try allocator.dupe(u8, buf[0..@intCast(n)]);
}
var ptr: [*]u8 = undefined;
const n = ext.kv_get(ns.ptr, ns.len, key.ptr, key.len, &ptr);
if (n < 0) return null;
if (n == 0) return try allocator.dupe(u8, "");
defer std.heap.wasm_allocator.free(ptr[0..@intCast(n)]);
return try allocator.dupe(u8, ptr[0..@intCast(n)]);
}
fn put(_: ?*anyopaque, ns: []const u8, key: []const u8, value: []const u8, opts: Kv.PutOptions) !void {
@ -36,10 +25,13 @@ fn delete(_: ?*anyopaque, ns: []const u8, key: []const u8) !void {
}
fn list(_: ?*anyopaque, ns: []const u8, allocator: std.mem.Allocator, prefix: []const u8) ![][]u8 {
var buf: [65536]u8 = undefined;
const n = ext.kv_list(ns.ptr, ns.len, prefix.ptr, prefix.len, &buf, buf.len);
if (n <= 0) return &[_][]u8{};
const parsed = try std.json.parseFromSlice([][]const u8, allocator, buf[0..@intCast(n)], .{});
var ptr: [*]u8 = undefined;
const n = ext.kv_list(ns.ptr, ns.len, prefix.ptr, prefix.len, &ptr);
if (n < 0) return error.InvalidResponse;
if (n == 0) return &[_][]u8{};
defer std.heap.wasm_allocator.free(ptr[0..@intCast(n)]);
const parsed = try std.json.parseFromSlice([][]const u8, allocator, ptr[0..@intCast(n)], .{});
defer parsed.deinit();
const keys = try allocator.alloc([]u8, parsed.value.len);
for (parsed.value, 0..) |k, i| keys[i] = try allocator.dupe(u8, k);
@ -52,8 +44,7 @@ const ext = struct {
ns_len: usize,
key_ptr: [*]const u8,
key_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
pub extern "__zx_kv" fn kv_put(
@ -78,8 +69,7 @@ const ext = struct {
ns_len: usize,
prefix_ptr: [*]const u8,
prefix_len: usize,
buf_ptr: [*]u8,
buf_max: usize,
out_ptr: *[*]u8,
) i32;
};