From f5a028ff76b69a55c35f2161f796dbae8f49ab51 Mon Sep 17 00:00:00 2001
From: Trevor Elliott <telliott@fastly.com>
Date: Thu, 5 Oct 2023 17:22:52 -0700
Subject: [PATCH] Update for pooling allocator and wasi-nn changes

---
 lib/src/execute.rs | 13 +++++++------
 lib/src/linking.rs | 17 ++++++++++-------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/lib/src/execute.rs b/lib/src/execute.rs
index 1546055c..df731fee 100644
--- a/lib/src/execute.rs
+++ b/lib/src/execute.rs
@@ -517,21 +517,22 @@ fn configure_wasmtime(profiling_strategy: ProfilingStrategy) -> wasmtime::Config
     let mut pooling_allocation_config = PoolingAllocationConfig::default();
 
     // This number matches C@E production
-    pooling_allocation_config.instance_size(MB);
+    pooling_allocation_config.max_core_instance_size(MB);
 
     // Core wasm programs have 1 memory
-    pooling_allocation_config.instance_memories(1);
+    pooling_allocation_config.total_memories(1);
+    pooling_allocation_config.max_memories_per_module(1);
 
     // allow for up to 128MiB of linear memory. Wasm pages are 64k
-    pooling_allocation_config.instance_memory_pages(128 * (MB as u64) / (64 * 1024));
+    pooling_allocation_config.memory_pages(128 * (MB as u64) / (64 * 1024));
 
     // Core wasm programs have 1 table
-    pooling_allocation_config.instance_tables(1);
+    pooling_allocation_config.max_tables_per_module(1);
 
     // Some applications create a large number of functions, in particular
     // when compiled in debug mode or applications written in swift. Every
     // function can end up in the table
-    pooling_allocation_config.instance_table_elements(98765);
+    pooling_allocation_config.table_elements(98765);
 
     // Maximum number of slots in the pooling allocator to keep "warm", or those
     // to keep around to possibly satisfy an affine allocation request or an
@@ -542,7 +543,7 @@ fn configure_wasmtime(profiling_strategy: ProfilingStrategy) -> wasmtime::Config
     // memory space if multiple engines are spun up in a single process. We'll likely want to move
     // to the on-demand allocator eventually for most purposes; see
     // https://github.com/fastly/Viceroy/issues/255
-    pooling_allocation_config.instance_count(100);
+    pooling_allocation_config.total_core_instances(100);
 
     config.allocation_strategy(InstanceAllocationStrategy::Pooling(
         pooling_allocation_config,
diff --git a/lib/src/linking.rs b/lib/src/linking.rs
index 496efefb..e768c300 100644
--- a/lib/src/linking.rs
+++ b/lib/src/linking.rs
@@ -89,7 +89,8 @@ pub(crate) fn create_store(
     guest_profiler: Option<GuestProfiler>,
 ) -> Result<Store<WasmCtx>, anyhow::Error> {
     let wasi = make_wasi_ctx(ctx, &session).context("creating Wasi context")?;
-    let wasi_nn = WasiNnCtx::new().unwrap();
+    let (backends, registry) = wasmtime_wasi_nn::preload(&[])?;
+    let wasi_nn = WasiNnCtx::new(backends, registry);
     let wasm_ctx = WasmCtx {
         wasi,
         wasi_nn,
@@ -117,22 +118,22 @@ fn make_wasi_ctx(ctx: &ExecuteCtx, session: &Session) -> Result<WasiCtx, anyhow:
     // Viceroy provides a subset of the `FASTLY_*` environment variables that the production
     // Compute@Edge platform provides:
 
-    wasi_ctx = wasi_ctx
+    wasi_ctx
         // signal that we're in a local testing environment
         .env("FASTLY_HOSTNAME", "localhost")?
         // request IDs start at 0 and increment, rather than being UUIDs, for ease of testing
         .env("FASTLY_TRACE_ID", &format!("{:032x}", session.req_id()))?;
 
     if ctx.log_stdout() {
-        wasi_ctx = wasi_ctx.stdout(Box::new(WritePipe::new(LogEndpoint::new(b"stdout"))));
+        wasi_ctx.stdout(Box::new(WritePipe::new(LogEndpoint::new(b"stdout"))));
     } else {
-        wasi_ctx = wasi_ctx.inherit_stdout();
+        wasi_ctx.inherit_stdout();
     }
 
     if ctx.log_stderr() {
-        wasi_ctx = wasi_ctx.stderr(Box::new(WritePipe::new(LogEndpoint::new(b"stderr"))));
+        wasi_ctx.stderr(Box::new(WritePipe::new(LogEndpoint::new(b"stderr"))));
     } else {
-        wasi_ctx = wasi_ctx.inherit_stderr();
+        wasi_ctx.inherit_stderr();
     }
     Ok(wasi_ctx.build())
 }
@@ -144,7 +145,9 @@ pub fn link_host_functions(
     experimental_modules
         .iter()
         .try_for_each(|experimental_module| match experimental_module {
-            ExperimentalModule::WasiNn => wasmtime_wasi_nn::add_to_linker(linker, WasmCtx::wasi_nn),
+            ExperimentalModule::WasiNn => {
+                wasmtime_wasi_nn::witx::add_to_linker(linker, WasmCtx::wasi_nn)
+            }
         })?;
     wasmtime_wasi::add_to_linker(linker, WasmCtx::wasi)?;
     wiggle_abi::fastly_abi::add_to_linker(linker, WasmCtx::session)?;