feat(cu): mem reset in eval

Author: VinceJuliano

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

@@ -12,11 +12,13 @@ export function evaluateWith ({
    * is passed in. Eventually remove usage and injection
    */
   loadWasmModule,
+  locateScheduler,
   wasmInstanceCache,
   bootstrapWasmInstance,
   saveEvaluation,
   addExtension,
   ARWEAVE_URL,
+  ENABLE_MEMORY_RESET,
   logger
 }) {
   loadWasmModule = fromPromise(loadWasmModule)
@@ -94,7 +96,9 @@ export function evaluateWith ({
        */
       Memory: ifElse(
         pathOr(undefined, ['Error']),
-        always(prevMemory),
+        (res) => {
+          return prevMemory
+        },
         (res) => {
           const output = cond([
             [is(String), identity],
@@ -128,7 +132,6 @@ export function evaluateWith ({
       Messages: pathOr([], ['Messages']),
       Assignments: pathOr([], ['Assignments']),
       Spawns: pathOr([], ['Spawns']),
-      Patches: pathOr([], ['Patches']),
       Output: pipe(
         pathOr('', ['Output']),
         /**
@@ -146,6 +149,48 @@ export function evaluateWith ({
     })
   )
 
+  async function evalInitialMemory ({ wasmInstance, AoGlobal }) {
+    logger('Initializing fresh memory for process "%s"', AoGlobal.Process.Id)
+
+    const scheduler = await locateScheduler(
+      AoGlobal.Process.Id,
+      '_GQ33BkPtZrqxA84vM8Zk-N2aO0toNNu_C-l-rawrBA'
+    )
+
+    const processMessageFetch = await fetch(
+      `${scheduler.url}/${AoGlobal.Process.Id}?limit=1`
+    ).then(res => res.json())
+    const processMessage = processMessageFetch.edges[0].node
+
+    const initMessage = {
+      Id: AoGlobal.Process.Id,
+      Signature: processMessage.message.signature,
+      Data: processMessage.message.data,
+      Owner: processMessage.message.owner.address,
+      Target: AoGlobal.Process.Id,
+      Anchor: processMessage.message.anchor,
+      From: processMessage.message.owner.address,
+      'Forwarded-By': undefined,
+      Tags: processMessage.message.tags,
+      Epoch: 0,
+      Nonce: 0,
+      Timestamp: parseInt(processMessage.assignment.tags.find(
+        (t) => t.name === 'Timestamp'
+      ).value),
+      'Block-Height': parseInt(processMessage.assignment.tags.find(
+        (t) => t.name === 'Block-Height'
+      ).value),
+      'Hash-Chain': processMessage.assignment.tags.find(
+        (t) => t.name === 'Hash-Chain'
+      ).value,
+      Cron: false,
+      'Read-Only': false
+    }
+    const mem = new Uint8Array(null)
+    const result = await wasmInstance(mem, initMessage, AoGlobal)
+    return result
+  }
+
   /**
    * Evaluate a message using the handler that wraps the WebAssembly.Instance,
    * identified by the streamId.
@@ -164,22 +209,39 @@ export function evaluateWith ({
    * Finally, evaluates the message and returns the result of the evaluation.
    */
   return ({ streamId, moduleId, wasmModule, moduleOptions, processId, noSave, name, deepHash, cron, ordinate, isAssignment, Memory, message, AoGlobal }) => {
+    const shouldReset = ENABLE_MEMORY_RESET && message.Tags && message.Tags.some(t =>
+      t.name === 'Action' &&
+      t.value === 'Reset-Memory' &&
+      AoGlobal.Process.Owner === message.Owner
+    )
+
+    if (shouldReset) {
+      wasmInstanceCache.delete(streamId)
+    }
+
     /**
      * Dynamically load the module, either from cache,
      * or from a file
      */
     return maybeCachedInstance({ streamId, moduleId, wasmModule, moduleOptions, name, processId, Memory, message, AoGlobal })
       .bichain(loadInstance, Resolved)
+      .chain(fromPromise(async (wasmInstance) => {
+        const maybeSwappedResult = shouldReset
+          ? await evalInitialMemory({ wasmInstance, message, AoGlobal })
+          : null
+
+        return { wasmInstance, Memory, maybeSwappedResult }
+      }))
       /**
        * Perform the evaluation
        */
-      .chain((wasmInstance) =>
-        of(wasmInstance)
-          .map((wasmInstance) => {
+      .chain(({ wasmInstance, Memory, maybeSwappedResult }) =>
+        of({ wasmInstance, Memory, maybeSwappedResult })
+          .map(({ wasmInstance, Memory, maybeSwappedResult }) => {
             logger('Evaluating message "%s" to process "%s"', name, processId)
-            return wasmInstance
+            return { wasmInstance, Memory, maybeSwappedResult }
           })
-          .chain(fromPromise(async (wasmInstance) =>
+          .chain(fromPromise(async ({ wasmInstance, Memory, maybeSwappedResult }) => {
             /**
              * AoLoader requires Memory to be a View, so that it can set the WebAssembly.Instance
              * memory.
@@ -196,8 +258,12 @@ export function evaluateWith ({
              *
              * TODO: check if any performance implications in using set() within AoLoaderrs
              */
-            wasmInstance(ArrayBuffer.isView(Memory) ? Memory : new Uint8Array(Memory), message, AoGlobal)
-          ))
+            if (maybeSwappedResult) {
+              return maybeSwappedResult
+            }
+            const mem = ArrayBuffer.isView(Memory) ? Memory : new Uint8Array(Memory)
+            return wasmInstance(mem, message, AoGlobal)
+          }))
           .bichain(
             /**
              * Map thrown error to a result.error. In this way, the Worker should _never_

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

@@ -1,3 +1,5 @@
+import { connect } from '@permaweb/ao-scheduler-utils'
+
 import * as WasmClient from '../../wasm.js'
 import * as AoEvaluationClient from '../../ao-evaluation.js'
 import * as DbClient from '../../db.js'
@@ -7,6 +9,13 @@ import { evaluateWith } from '../evaluate.js'
 export const createApis = async (ctx) => {
   const db = await DbClient.createDbClient({ url: ctx.DB_URL, bootstrap: false, max: 5 })
   const wasmInstanceCache = WasmClient.createWasmInstanceCache({ MAX_SIZE: ctx.WASM_INSTANCE_CACHE_MAX_SIZE })
+  const { locate } = connect({
+    cacheSize: 100,
+    GRAPHQL_URL: ctx.GRAPHQL_URL,
+    followRedirects: true,
+    GRAPHQL_MAX_RETRIES: 5,
+    GRAPHQL_RETRY_BACKOFF: 300
+  })
 
   const close = async (streamId) => wasmInstanceCache.delete(streamId)
 
@@ -24,7 +33,9 @@ export const createApis = async (ctx) => {
       logger: ctx.logger
     }),
     ARWEAVE_URL: ctx.ARWEAVE_URL,
-    logger: ctx.logger
+    locateScheduler: locate,
+    logger: ctx.logger,
+    ENABLE_MEMORY_RESET: true
   })
 
   return { evaluate, close }