ui: write wasm RPC requests directly into the tokenizer

gn/standalone/wasm.gni

@@ -47,11 +47,12 @@ template("wasm_lib") {
       _environment = "worker"
     }
     if (_enable_fs) {
-      _exports = "['ccall', 'callMain', 'addFunction', 'FS', 'FS_mkdir', " +
+      _exports = "['ccall', 'cwrap', 'callMain', 'addFunction', 'FS', " +
+                 "'FS_mkdir', " +
                  "'FS_mount', 'FS_lookupPath', 'FS_unlink', 'FS_readdir', " +
                  "'FS_readFile', 'WORKERFS', 'HEAPU8']"
     } else {
-      _exports = "['ccall', 'callMain', 'addFunction', 'HEAPU8']"
+      _exports = "['ccall', 'cwrap', 'callMain', 'addFunction', 'HEAPU8']"
     }
     _target_ldflags = [
       "-s",

gn/standalone/wasm_typescript_declaration.d.ts

@@ -48,7 +48,7 @@ declare namespace Wasm {
 
   export interface Module {
     callMain(args: string[]): void;
-    addFunction(f: any, argTypes: string): void;
+    addFunction(f: any, argTypes: string): number;
     FS_mkdir(path: string, mode?: number): any;
     FS_mount(type: Wasm.FileSystemType, opts: any, mountpoint: string): any;
     FS_lookupPath(path: string): { path: string; node: Wasm.FileSystemNode };
@@ -63,6 +63,12 @@ declare namespace Wasm {
       argTypes: string[],
       args: any[],
     ): number;
+    // Like ccall(), but resolves the symbol and marshalling once.
+    cwrap(
+      ident: string,
+      returnType: string,
+      argTypes: string[],
+    ): (...args: number[]) => number;
     HEAPU8: Uint8Array;
     FS: FileSystem;
   }

src/trace_processor/rpc/wasm_bridge.cc

@@ -30,11 +30,9 @@ namespace {
 using RpcResponseFn = void(const void*, uint32_t);
 
 Rpc* g_trace_processor_rpc;
+Rpc::RequestHandle* g_pending_request;
 
-// The buffer used to pass the request arguments. The caller (JS) decides how
-// big this buffer should be in the Initialize() call.
-uint8_t* g_req_buf;
-uint32_t g_req_buf_size;
+uint32_t g_max_write_size;
 
 PERFETTO_NO_INLINE void OutOfMemoryHandler() {
   fprintf(stderr, "\nCannot enlarge memory\n");
@@ -48,11 +46,11 @@ PERFETTO_NO_INLINE void OutOfMemoryHandler() {
 // +---------------------------------------------------------------------------+
 extern "C" {
 
-// Returns the address of the allocated request buffer.
+// Returns the address JS must write the first request into.
 uint8_t* EMSCRIPTEN_KEEPALIVE
 trace_processor_rpc_init(RpcResponseFn* RpcResponseFn, uint32_t);
 uint8_t* trace_processor_rpc_init(RpcResponseFn* resp_function,
-                                  uint32_t req_buffer_size) {
+                                  uint32_t max_write_size) {
   // Usually OOMs manifest as a failure in dlmalloc() -> sbrk() ->
   //_emscripten_resize_heap() which aborts itself. However in some rare cases
   // sbrk() can fail outside of _emscripten_resize_heap and just return null.
@@ -69,20 +67,19 @@ uint8_t* trace_processor_rpc_init(RpcResponseFn* resp_function,
   // an overview of the JS<>Wasm callstack.
   g_trace_processor_rpc->SetRpcResponseFunction(resp_function);
 
-  g_req_buf = new uint8_t[req_buffer_size];
-  g_req_buf_size = req_buffer_size;
-  return g_req_buf;
+  g_max_write_size = max_write_size;
+  g_pending_request = new Rpc::RequestHandle(
+      g_trace_processor_rpc->BeginRpcRequest(g_max_write_size));
+  return g_pending_request->data();
 }
 
-void EMSCRIPTEN_KEEPALIVE trace_processor_on_rpc_request(uint32_t);
-void trace_processor_on_rpc_request(uint32_t size) {
-  if (PERFETTO_UNLIKELY(size > g_req_buf_size)) {
-    fprintf(stderr,
-            "RPC request size exceeds the buffer passed to "
-            "trace_processor_rpc_init\n");
-    return;
-  }
-  g_trace_processor_rpc->OnRpcRequest(g_req_buf, size);
+// Returning the next address rather than taking one keeps this to a single
+// JS->Wasm call per request: JS knows where to write before it knows the size.
+uint8_t* EMSCRIPTEN_KEEPALIVE trace_processor_on_rpc_request(uint32_t);
+uint8_t* trace_processor_on_rpc_request(uint32_t size) {
+  g_pending_request->EndRequest(size);
+  *g_pending_request = g_trace_processor_rpc->BeginRpcRequest(g_max_write_size);
+  return g_pending_request->data();
 }
 
 }  // extern "C"

ui/src/base/proto_utils_wasm.ts

@@ -24,9 +24,20 @@ import WasmModuleGen from '../gen/proto_utils';
  * It guarantees to have the same behaviour of perfetto_cmd and trace_processor
  * by using precisely the same code via WebAssembly.
  */
+// Each takes the number of bytes in |buf| and returns the number written back.
+const kConvertFns = [
+  'trace_config_pb_to_txt',
+  'trace_config_txt_to_pb',
+  'trace_summary_spec_to_text',
+  'trace_summary_spec_txt_to_pb',
+] as const;
+type ConvertFn = (typeof kConvertFns)[number];
+
 interface WasmModule {
   module: WasmModuleGen.Module;
   buf: Uint8Array;
+  // cwrap() resolves the export once; ccall() would redo it on every call.
+  fns: Record<ConvertFn, (size: number) => number>;
 }
 
 let moduleInstance: WasmModule | undefined = undefined;
@@ -38,7 +49,7 @@ let moduleInstance: WasmModule | undefined = undefined;
 async function pbToText<T>(
   input: Uint8Array | T,
   encode: (val: T) => {finish(): Uint8Array},
-  ccallName: string,
+  fnName: ConvertFn,
 ): Promise<string> {
   const wasm = await initWasmOnce();
   const inputU8: Uint8Array =
@@ -50,8 +61,7 @@ async function pbToText<T>(
   }
   wasm.buf.set(inputU8);
 
-  const txtSize =
-    wasm.module.ccall(ccallName, 'number', ['number'], [inputU8.length]) >>> 0;
+  const txtSize = wasm.fns[fnName](inputU8.length) >>> 0;
   return utf8Decode(wasm.buf.subarray(0, txtSize));
 }
 
@@ -71,7 +81,7 @@ export async function traceConfigToTxt(
  */
 async function textToPb(
   input: string,
-  ccallName: string,
+  fnName: ConvertFn,
 ): Promise<Result<Uint8Array>> {
   const wasm = await initWasmOnce();
 
@@ -83,9 +93,7 @@ async function textToPb(
   }
   wasm.buf.set(inputUtf8);
 
-  const resSize =
-    wasm.module.ccall(ccallName, 'number', ['number'], [inputUtf8.length]) >>>
-    0;
+  const resSize = wasm.fns[fnName](inputUtf8.length) >>> 0;
 
   const success = wasm.buf.at(0) === 1;
   const payload = wasm.buf.slice(1, 1 + resSize);
@@ -133,12 +141,17 @@ async function initWasmOnce(): Promise<WasmModule> {
       onRuntimeInitialized: () => {},
       wasmBinary,
     } as WasmModuleGen.ModuleArgs);
-    const bufAddr = instance.ccall('proto_utils_buf', 'number', [], []) >>> 0;
+    const bufAddr = instance.cwrap('proto_utils_buf', 'number', [])() >>> 0;
     const bufSize =
-      instance.ccall('proto_utils_buf_size', 'number', [], []) >>> 0;
+      instance.cwrap('proto_utils_buf_size', 'number', [])() >>> 0;
+    const fns = {} as Record<ConvertFn, (size: number) => number>;
+    for (const name of kConvertFns) {
+      fns[name] = instance.cwrap(name, 'number', ['number']);
+    }
     moduleInstance = {
       module: instance,
       buf: instance.HEAPU8.subarray(bufAddr, bufAddr + bufSize),
+      fns,
     };
   }
   return moduleInstance;

ui/src/engine/wasm_bridge.ts

@@ -19,16 +19,16 @@ import {
   memory64Supported,
 } from '../trace_processor/wasm_modules';
 
-// The Initialize() call will allocate a buffer of REQ_BUF_SIZE bytes which
-// will be used to copy the input request data. This is to avoid passing the
-// input data on the stack, which has a limited (~1MB) size.
-// The buffer will be allocated by the C++ side and reachable at
-// HEAPU8[reqBufferAddr, +REQ_BUFFER_SIZE].
-const REQ_BUF_SIZE = 32 * 1024 * 1024;
+// Also the size of the standing reservation held for us in the C++ tokenizer.
+// Must be <= ProtoRingBuffer::kMaxMsgSize (64MB), and is kept just above
+// TRACE_SLICE_SIZE (32MB, trace_stream.ts) so a whole TPM_APPEND_TRACE_DATA
+// lands in one write; if that stops holding, writes are simply split.
+const MAX_WRITE_SIZE = 33 * 1024 * 1024;
 
 // The end-to-end interaction between JS and Wasm is as follows:
 // - [JS] Inbound data received by the worker (onmessage() in engine/index.ts).
-//   - [JS] onRpcDataReceived() (this file)
+//   - [JS] onRpcDataReceived() (this file) writes the bytes straight into the
+//     address that the previous call handed back, then
 //     - [C++] trace_processor_on_rpc_request (wasm_bridge.cc)
 //       - [C++] some TraceProcessor::method()
 //         for (batch in result_rows)
@@ -38,7 +38,10 @@ const REQ_BUF_SIZE = 32 * 1024 * 1024;
 export class WasmBridge {
   private aborted = false;
   private connection?: TraceProcessor64.Module;
-  private reqBufferAddr = 0;
+  // cwrap() resolves the export once; ccall() would redo it on every call.
+  private onRpcRequest?: (size: number) => number;
+  // Where the next request goes; each call hands back the one after it.
+  private wrAddr = 0;
   private lastStderr: string[] = [];
   private messagePort?: MessagePort;
   private useMemory64 = false;
@@ -68,13 +71,16 @@ export class WasmBridge {
       },
     });
     const fn = connection.addFunction(this.onReply.bind(this), 'vpi');
-    this.reqBufferAddr = this.wasmPtrCast(
-      connection.ccall(
-        'trace_processor_rpc_init',
-        /* return=*/ 'pointer',
-        /* args=*/ ['pointer', 'number'],
-        [fn, REQ_BUF_SIZE],
-      ),
+    const init = connection.cwrap(
+      'trace_processor_rpc_init',
+      /* return=*/ 'pointer',
+      /* args=*/ ['pointer', 'number'],
+    );
+    this.wrAddr = this.wasmPtrCast(init(fn, MAX_WRITE_SIZE));
+    this.onRpcRequest = connection.cwrap(
+      'trace_processor_on_rpc_request',
+      /* return=*/ 'pointer',
+      /* args=*/ ['number'],
     );
     this.connection = connection;
 
@@ -91,20 +97,17 @@ export class WasmBridge {
     assertTrue(msg.data instanceof Uint8Array);
     const data = msg.data as Uint8Array;
     let wrSize = 0;
-    // If the request data is larger than our JS<>Wasm interop buffer, split it
-    // into multiple writes. The RPC channel is byte-oriented and is designed to
-    // deal with arbitrary fragmentations.
+    // If the request data is larger than MAX_WRITE_SIZE, split it into multiple
+    // writes. The RPC channel is byte-oriented and is designed to deal with
+    // arbitrary fragmentations.
     while (wrSize < data.length) {
-      const sliceLen = Math.min(data.length - wrSize, REQ_BUF_SIZE);
+      const sliceLen = Math.min(data.length - wrSize, MAX_WRITE_SIZE);
       const dataSlice = data.subarray(wrSize, wrSize + sliceLen);
-      connection.HEAPU8.set(dataSlice, this.reqBufferAddr);
+      connection.HEAPU8.set(dataSlice, this.wrAddr);
       wrSize += sliceLen;
       try {
-        connection.ccall(
-          'trace_processor_on_rpc_request', // C function name.
-          'void', // Return type.
-          ['number'], // Arg types.
-          [sliceLen], // Args.
+        this.wrAddr = this.wasmPtrCast(
+          ensureExists(this.onRpcRequest)(sliceLen),
         );
       } catch (err) {
         this.aborted = true;
@@ -119,7 +122,7 @@ export class WasmBridge {
   }
 
   // This function is bound and passed to Initialize and is called by the C++
-  // code while in the ccall(trace_processor_on_rpc_request).
+  // code while in the call to trace_processor_on_rpc_request.
   private onReply(heapPtrArg: bigint | number, size: number) {
     const heapPtr = this.wasmPtrCast(heapPtrArg);
     const data = ensureExists(this.connection).HEAPU8.slice(