fix(cu): transfer memory between worker and main thread

servers/cu/src/domain/client/ao-module.js

@@ -4,8 +4,12 @@ import { fromPromise, of, Rejected, Resolved } from 'hyper-async'
 import { always, applySpec, defaultTo, evolve, head, prop } from 'ramda'
 import { z } from 'zod'
 
+import { arrayBufferFromMaybeView } from '../utils.js'
 import { moduleSchema } from '../model.js'
 import { MODULES_TABLE } from './sqlite.js'
+import { timer } from './metrics.js'
+
+const TWO_GB = 2 * 1024 * 1024 * 1024
 
 const moduleDocSchema = z.object({
   id: z.string().min(1),
@@ -116,6 +120,69 @@ export function evaluatorWith ({ evaluate, loadWasmModule }) {
                */
               if (defer) await new Promise(resolve => setImmediate(resolve))
 
+              if (args.Memory) {
+                /**
+                 * The ArrayBuffer is transferred to the worker as part of performing
+                 * an evaluation. This transfer will subsequently detach any views, Buffers,
+                 * and more broadly, references to the ArrayBuffer on this thread.
+                 *
+                 * So if this is the first eval being performed for the eval stream,
+                 * then we copy the contents of the ArrayBuffer. That way, we can be sure
+                 * that no references on the main thread will be affected during the eval stream
+                 * transfers happening back and forth. This effectively give's each eval stream
+                 * it's own ArrayBuffer to pass back and forth.
+                 *
+                 * (this is no worse than the structured clone that was happening before
+                 * as part of message passing. But instead, the clone is only performed once,
+                 * instead of on each evaluation)
+                 *
+                 * TODO: perhaps there is a way to somehow lock the ArrayBuffer usage
+                 * instead of copying on first evaluation. We have to be careful that nothing
+                 * (ie. a view of the ArrayBuffer in a Wasm Instnace dryrun)
+                 * inadvertantly mutates the underlying ArrayBuffer
+                 */
+                if (args.first) {
+                  let stopTimer = () => {}
+                  if (args.Memory.byteLength > TWO_GB) {
+                    stopTimer = timer('copyLargeMemory', {
+                      streamId,
+                      processId: args.processId,
+                      byteLength: args.Memory.byteLength
+                    }).stop
+                  }
+                  /**
+                   * We must pass a view into copyBytesFrom,
+                   *
+                   * so we first check whether it already is or not,
+                   * and create one on top of the ArrayBuffer if necessary
+                   *
+                   * (NodeJS' Buffer is a subclass of DataView)
+                   */
+                  args.Memory = ArrayBuffer.isView(args.Memory)
+                    ? Buffer.copyBytesFrom(args.Memory)
+                    : Buffer.copyBytesFrom(new Uint8Array(args.Memory))
+                  stopTimer()
+                }
+
+                /**
+                 * If Memory is sufficiently large, transferring the View somehow
+                 * causes the underlying ArrayBuffer to be truncated. This truncation
+                 * does not occur when instead the underlying ArrayBuffer is transferred,
+                 * directly.
+                 *
+                 * So we always ensure the Memory transferred to the worker thread
+                 * is the actual ArrayBuffer, and not a View.
+                 *
+                 * (the same is done in the opposite direction in the worker thread)
+                 *
+                 * TODO: maybe AoLoader should be made to return the underlying ArrayBuffer
+                 * as Memory, instead of a View?
+                 */
+                args.Memory = arrayBufferFromMaybeView(args.Memory)
+
+                options = { transfer: [args.Memory] }
+              }
+
               args.streamId = streamId
               args.moduleId = moduleId
               args.moduleOptions = moduleOptions

servers/cu/src/domain/lib/evaluate.js

@@ -157,6 +157,7 @@ export function evaluateWith (env) {
                  * Iterate over the async iterable of messages,
                  * and evaluate each one
                  */
+                let first = true
                 for await (const { noSave, cron, ordinate, name, message, deepHash, isAssignment, AoGlobal } of messages) {
                   if (cron) {
                     const key = toEvaledCron({ timestamp: message.Timestamp, cron })
@@ -208,14 +209,19 @@ export function evaluateWith (env) {
                         /**
                          * Where the actual evaluation is performed
                          */
-                        .then((Memory) => ctx.evaluator({ noSave, name, deepHash, cron, ordinate, isAssignment, processId: ctx.id, Memory, message, AoGlobal }))
+                        .then((Memory) => ctx.evaluator({ first, noSave, name, deepHash, cron, ordinate, isAssignment, processId: ctx.id, Memory, message, AoGlobal }))
                         /**
                          * These values are folded,
                          * so that we can potentially update the process memory cache
                          * at the end of evaluation
                          */
                         .then(mergeLeft({ noSave, message, cron, ordinate }))
                         .then(async (output) => {
+                          /**
+                           * Make sure to set first to false
+                           * for all subsequent evaluations for this evaluation stream
+                           */
+                          if (first) first = false
                           if (output.GasUsed) totalGasUsed += BigInt(output.GasUsed ?? 0)
 
                           if (cron) ctx.stats.messages.cron++

servers/cu/src/domain/utils.js

@@ -426,3 +426,7 @@ export const maybeParseInt = pipe(
   mapBadData,
   (val) => val ? parseInt(val) : undefined
 )
+
+export const arrayBufferFromMaybeView = (maybeView) => ArrayBuffer.isView(maybeView)
+  ? maybeView.buffer
+  : maybeView // assumes an ArrayBuffer. TODO: maybe an additional check

servers/cu/src/domain/worker/evaluate.js

@@ -7,6 +7,10 @@ import { saveEvaluationSchema } from '../dal.js'
 const WASM_64_FORMAT = 'wasm64-unknown-emscripten-draft_2024_02_15'
 
 export function evaluateWith ({
+  /**
+   * TODO: no longer needed since the wasmModule
+   * is passed in. Eventually remove usage and injection
+   */
   loadWasmModule,
   wasmInstanceCache,
   bootstrapWasmInstance,
@@ -165,7 +169,7 @@ export function evaluateWith ({
      */
     if (close) {
       wasmInstanceCache.delete(streamId)
-      return
+      return Promise.resolve()
     }
     /**
      * Dynamically load the module, either from cache,
@@ -182,7 +186,25 @@ export function evaluateWith ({
             logger('Evaluating message "%s" to process "%s"', name, processId)
             return wasmInstance
           })
-          .chain(fromPromise(async (wasmInstance) => wasmInstance(Memory, message, AoGlobal)))
+          .chain(fromPromise(async (wasmInstance) =>
+            /**
+             * AoLoader requires Memory to be a View, so that it can set the WebAssembly.Instance
+             * memory.
+             *
+             * DataView.set is used internally in AoLoader to "reload" the Memory
+             * into a WebAssembly.Instance. set() claims to
+             * "intelligently copy the source range of the buffer to the destination range
+             * [if they share the same underlying ArrayBuffer]" but doesn't go into detail about what
+             * "intelligent" actually means.
+             *
+             * It seems to imply that only diffs will be copied, if any at all. So for the same
+             * exact underlying ArrayBuffer with no diffs, I would _hope_ that is little to no
+             * performance penalty, but we should verify
+             *
+             * TODO: check if any performance implications in using set() within AoLoaderrs
+             */
+            wasmInstance(ArrayBuffer.isView(Memory) ? Memory : new Uint8Array(Memory), message, AoGlobal)
+          ))
           .bichain(
             /**
              * Map thrown error to a result.error. In this way, the Worker should _never_

servers/cu/src/domain/worker/evaluator/index.js

@@ -2,9 +2,10 @@ import { workerData } from 'node:worker_threads'
 import { hostname } from 'node:os'
 
 import { fetch, setGlobalDispatcher, Agent } from 'undici'
-import { worker } from 'workerpool'
+import { worker, Transfer } from 'workerpool'
 
 import { createLogger } from '../../logger.js'
+import { arrayBufferFromMaybeView } from '../../utils.js'
 
 import { createApis } from './main.js'
 
@@ -26,4 +27,43 @@ const apis = await createApis({
 /**
  * Expose our worker api
  */
-worker(apis)
+worker({
+  evaluate: (...args) => apis.evaluate(...args)
+    /**
+     * Transfer the ownership of the underlying ArrayBuffer back to the main thread
+     * to prevent copying it over
+     */
+    .then((output) => {
+      /**
+       * The evaluation stream is being closed,
+       * so no output is returned, so nothing
+       * needs to be transferred
+       */
+      if (!output) return output
+      /**
+       * If the very first evaluation produces
+       * an error, the resultant Memory will be null
+       * (prevMemory is used, which initializes as null)
+       *
+       * So in this edge-case, there's nothing to transfer,
+       * so we simply return output
+       */
+      if (!output.Memory) return output
+
+      /**
+       * If Memory is sufficiently large, transferring the View somehow
+       * causes the underlying ArrayBuffer to be truncated. This truncation
+       * does not occur when instead the underlying ArrayBuffer is transferred,
+       * directly.
+       *
+       * So we always ensure the Memory transferred back to the main thread
+       * is the actual ArrayBuffer, and not a View.
+       *
+       * TODO: maybe AoLoader should be made to return the underlying ArrayBuffer
+       * as Memory, instead of a View?
+       */
+      output.Memory = arrayBufferFromMaybeView(output.Memory)
+
+      return new Transfer(output, [output.Memory])
+    })
+})

servers/cu/src/domain/worker/evaluator/main.js

@@ -8,6 +8,10 @@ export const createApis = async (ctx) => {
   const sqlite = await SqliteClient.createSqliteClient({ url: ctx.DB_URL, bootstrap: false })
 
   const evaluate = evaluateWith({
+    /**
+     * TODO: no longer needed since the wasmModule
+     * is passed in. Eventually remove
+     */
     loadWasmModule: WasmClient.loadWasmModuleWith({
       fetch,
       ARWEAVE_URL: ctx.ARWEAVE_URL,