Merge pull request #153 from technosophos/feat/load-modules-on-startup

Load modules at startup, and clone per request

Author: technosophos

src/dispatcher.rs

@@ -1,11 +1,14 @@
 use std::net::SocketAddr;
+use std::path::Path;
+use std::sync::Arc;
 
 use hyper::{
     http::request::Parts,
     Body, Request, Response, StatusCode,
 };
 use sha2::{Digest, Sha256};
 use tracing::{instrument};
+use wasmtime::{Engine, Config};
 
 use crate::dynamic_route::{DynamicRoutes, interpret_routes};
 use crate::emplacer::Bits;
@@ -16,7 +19,7 @@ use crate::request::{RequestContext, RequestGlobalContext};
 
 use crate::bindle_util::{WagiHandlerInfo};
 use crate::wagi_config::{LoadedHandlerConfiguration, ModuleMapConfigurationEntry};
-use crate::wasm_module::WasmModuleSource;
+use crate::wasm_module::CompiledWasmModule;
 use crate::wasm_runner::{RunWasmResult, prepare_stdio_streams, prepare_wasm_instance, run_prepared_wasm_instance_if_present, WasmLinkOptions};
 
 #[derive(Clone, Debug)]
@@ -89,12 +92,33 @@ impl RoutingTableEntry {
         self.route_pattern.is_match(uri_fragment)
     }
 
-    fn build_from_modules_toml(source: &Loaded<ModuleMapConfigurationEntry>) -> anyhow::Result<RoutingTableEntry> {
+    /// Create a new Wasm Engine and configure it.
+    fn new_engine(cache_config_path: &Path) -> anyhow::Result<Engine> {
+        let mut config = Config::default();
+
+        // Enable multi memory and module linking support.
+        config.wasm_multi_memory(true);
+        config.wasm_module_linking(true);
+
+        if let Ok(p) = std::fs::canonicalize(cache_config_path) {
+            config.cache_config_load(p)?;
+        };
+
+        Engine::new(&config)
+    }
+
+    fn compile_module(data: Arc<Vec<u8>>, cache_config_path: &Path) -> anyhow::Result<CompiledWasmModule> {
+        let engine = RoutingTableEntry::new_engine(cache_config_path)?;
+        let module = wasmtime::Module::new(&engine, &**data)?;
+        Ok(CompiledWasmModule::Object(module, engine))
+    }
+
+    fn build_from_modules_toml(source: &Loaded<ModuleMapConfigurationEntry>, global_context: &RequestGlobalContext,) -> anyhow::Result<RoutingTableEntry> {
         let route_pattern = RoutePattern::parse(&source.metadata.route);
 
-        let wasm_source = WasmModuleSource::Blob(source.content.clone());
+        let wasm_module = RoutingTableEntry::compile_module(source.content.clone(), &global_context.cache_config_path)?;
         let wasm_route_handler = WasmRouteHandler {
-            wasm_module_source: wasm_source,
+            wasm_module_source: wasm_module,
             wasm_module_name: source.metadata.module.clone(),
             entrypoint: source.metadata.entrypoint.clone().unwrap_or_else(|| DEFAULT_ENTRYPOINT.to_owned()),
             volumes: source.metadata.volumes.clone().unwrap_or_default(),
@@ -109,25 +133,29 @@ impl RoutingTableEntry {
         })
     }
 
-    fn build_from_bindle_entry(source: &(WagiHandlerInfo, Bits)) -> Option<anyhow::Result<RoutingTableEntry>> {
+    fn build_from_bindle_entry(source: &(WagiHandlerInfo, Bits), global_context: &RequestGlobalContext,) -> Option<anyhow::Result<RoutingTableEntry>> {
         let (wagi_handler, bits) = source;
 
         let route_pattern = RoutePattern::parse(&wagi_handler.route);
-        let wasm_source = WasmModuleSource::Blob(bits.wasm_module.clone());
-        let wasm_route_handler = WasmRouteHandler {
-            wasm_module_source: wasm_source,
-            wasm_module_name: wagi_handler.parcel.label.name.clone(),
-            entrypoint: wagi_handler.entrypoint.clone().unwrap_or_else(|| DEFAULT_ENTRYPOINT.to_owned()),
-            volumes: bits.volume_mounts.clone(),
-            allowed_hosts: wagi_handler.allowed_hosts.clone(),
-            http_max_concurrency: None,
-        };
-        let handler_info = RouteHandler::Wasm(wasm_route_handler);
+        match RoutingTableEntry::compile_module(bits.wasm_module.clone(), &global_context.cache_config_path) {
+            Err(e) => Some(Err(e)), // Not clear what we are supposed to return here.
+            Ok(wasm_module) => {
+                let wasm_route_handler = WasmRouteHandler {
+                    wasm_module_source: wasm_module,
+                    wasm_module_name: wagi_handler.parcel.label.name.clone(),
+                    entrypoint: wagi_handler.entrypoint.clone().unwrap_or_else(|| DEFAULT_ENTRYPOINT.to_owned()),
+                    volumes: bits.volume_mounts.clone(),
+                    allowed_hosts: wagi_handler.allowed_hosts.clone(),
+                    http_max_concurrency: None,
+                };
+                let handler_info = RouteHandler::Wasm(wasm_route_handler);
 
-        Some(Ok(Self {
-            route_pattern,
-            handler_info,
-        }))
+                Some(Ok(Self {
+                    route_pattern,
+                    handler_info,
+                }))
+            }
+        }
     }
 
     fn inbuilt(path: &str, handler: RouteHandler) -> Self {
@@ -263,9 +291,9 @@ impl RoutingTable {
     pub fn build(source: &LoadedHandlerConfiguration, global_context: RequestGlobalContext) -> anyhow::Result<RoutingTable> {
         let user_entries = match source {
             LoadedHandlerConfiguration::ModuleMapFile(module_map_entries) =>
-                Self::build_from_modules_toml(module_map_entries),
+                Self::build_from_modules_toml(module_map_entries, &global_context),
             LoadedHandlerConfiguration::Bindle(bindle_entries) =>
-                Self::build_from_bindle_entries(bindle_entries),
+                Self::build_from_bindle_entries(bindle_entries, &global_context),
         }?;
         let full_user_entries = augment_dynamic_routes(user_entries, &global_context)?;
 
@@ -278,18 +306,18 @@ impl RoutingTable {
         })
     }
 
-    fn build_from_modules_toml(module_map_entries: &[Loaded<ModuleMapConfigurationEntry>]) -> anyhow::Result<Vec<RoutingTableEntry>> {
+    fn build_from_modules_toml(module_map_entries: &[Loaded<ModuleMapConfigurationEntry>], global_context: &RequestGlobalContext) -> anyhow::Result<Vec<RoutingTableEntry>> {
         // TODO: look for `_routes` function
         module_map_entries
             .iter()
-            .map(|e| RoutingTableEntry::build_from_modules_toml(e))
+            .map(|e| RoutingTableEntry::build_from_modules_toml(e, global_context))
             .collect()
     }
 
-    fn build_from_bindle_entries(bindle_entries: &[(WagiHandlerInfo, Bits)]) -> anyhow::Result<Vec<RoutingTableEntry>> {
+    fn build_from_bindle_entries(bindle_entries: &[(WagiHandlerInfo, Bits)], global_context: &RequestGlobalContext) -> anyhow::Result<Vec<RoutingTableEntry>> {
         bindle_entries
             .iter()
-            .filter_map(|e| RoutingTableEntry::build_from_bindle_entry(e))
+            .filter_map(|e| RoutingTableEntry::build_from_bindle_entry(e, global_context))
             .collect()
     }
 
@@ -318,7 +346,7 @@ fn augment_one_wasm_with_dynamic_routes(routing_table_entry: &RoutingTableEntry,
 
     let ctx = build_wasi_context_for_dynamic_route_query(redirects.streams);
     let link_options = WasmLinkOptions::none();
-    let (store, instance) = prepare_wasm_instance(global_context, ctx, &wasm_route_handler.wasm_module_source, link_options)?;
+    let (store, instance) = prepare_wasm_instance(ctx, &wasm_route_handler.wasm_module_source, link_options)?;
 
     match run_prepared_wasm_instance_if_present(instance, store, "_routes") {
         RunWasmResult::WasmError(e) => Err(e),

src/handlers.rs

@@ -16,7 +16,7 @@ use crate::dispatcher::RoutePattern;
 use crate::http_util::{internal_error, parse_cgi_headers};
 use crate::request::{RequestContext, RequestGlobalContext};
 
-use crate::wasm_module::WasmModuleSource;
+use crate::wasm_module::CompiledWasmModule;
 use crate::wasm_runner::{prepare_stdio_streams, prepare_wasm_instance, run_prepared_wasm_instance, WasmLinkOptions};
 
 #[derive(Clone, Debug)]
@@ -27,7 +27,7 @@ pub enum RouteHandler {
 
 #[derive(Clone, Debug)]
 pub struct WasmRouteHandler {
-    pub wasm_module_source: WasmModuleSource,
+    pub wasm_module_source: CompiledWasmModule,
     pub wasm_module_name: String,
     pub entrypoint: String,
     pub volumes: HashMap<String, String>,
@@ -60,7 +60,7 @@ impl WasmRouteHandler {
 
         let ctx = self.build_wasi_context_for_request(req, headers, redirects.streams)?;
 
-        let (store, instance) = self.prepare_wasm_instance(global_context, ctx)?;
+        let (store, instance) = self.prepare_wasm_instance(ctx)?;
 
         // Drop manually to get instantiation time
         drop(startup_span);
@@ -103,11 +103,11 @@ impl WasmRouteHandler {
         Ok(ctx)
     }
 
-    fn prepare_wasm_instance(&self, global_context: &RequestGlobalContext, ctx: WasiCtx) -> Result<(Store<WasiCtx>, Instance), Error> {
+    fn prepare_wasm_instance(&self,  ctx: WasiCtx) -> Result<(Store<WasiCtx>, Instance), Error> {
         debug!("Preparing Wasm instance.");
         let link_options = WasmLinkOptions::default()
             .with_http(self.allowed_hosts.clone(), self.http_max_concurrency);
-        prepare_wasm_instance(global_context, ctx, &self.wasm_module_source, link_options)
+        prepare_wasm_instance(ctx, &self.wasm_module_source, link_options)
     }
 }
 

src/wasm_module.rs

@@ -2,27 +2,18 @@ use std::{fmt::Debug, sync::{Arc, RwLock}};
 
 use wasi_common::pipe::{ReadPipe, WritePipe};
 use wasmtime::*;
-use wasmtime_wasi::*;
 
 // In future this might be pre-instantiated or something like that, so we will
 // just abstract it to be safe.
 #[derive(Clone)]
-pub enum WasmModuleSource {
-    Blob(Arc<Vec<u8>>),
+pub enum CompiledWasmModule {
+    Object(Module, Engine)
 }
 
-impl WasmModuleSource {
-    pub fn load_module(&self, store: &Store<WasiCtx>) -> anyhow::Result<wasmtime::Module> {
-        match self {
-            Self::Blob(bytes) => wasmtime::Module::new(store.engine(), &**bytes),
-        }
-    }
-}
-
-impl Debug for WasmModuleSource {
+impl Debug for CompiledWasmModule {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::Blob(v) => f.write_fmt(format_args!("Blob(length={})", v.len())),
+            Self::Object(m, _) => f.write_fmt(format_args!("Object(Module={:?})", m.name())),
         }
     }
 }

src/wasm_runner.rs

@@ -1,4 +1,3 @@
-use std::path::Path;
 use std::sync::{Arc, RwLock};
 
 use wasi_common::pipe::{ReadPipe, WritePipe};
@@ -8,7 +7,7 @@ use wasmtime_wasi::*;
 use tracing::debug;
 
 use crate::request::RequestGlobalContext;
-use crate::wasm_module::WasmModuleSource;
+use crate::wasm_module::CompiledWasmModule;
 
 const STDERR_FILE: &str = "module.stderr";
 
@@ -80,39 +79,28 @@ pub fn prepare_stdio_streams(
     })
 }
 
-pub fn new_store_and_engine(
-    cache_config_path: &Path,
-    ctx: WasiCtx,
-) -> Result<(Store<WasiCtx>, Engine), anyhow::Error> {
-    let mut config = Config::default();
-
-    // Enable multi memory and module linking support.
-    config.wasm_multi_memory(true);
-    config.wasm_module_linking(true);
 
-    if let Ok(p) = std::fs::canonicalize(cache_config_path) {
-        config.cache_config_load(p)?;
-    };
-
-    let engine = Engine::new(&config)?;
-    Ok((Store::new(&engine, ctx), engine))
+pub fn new_store(
+    ctx: WasiCtx,
+    engine: &Engine
+) -> Result<Store<WasiCtx>, anyhow::Error> {
+    Ok(Store::new(engine, ctx))
 }
 
 pub fn prepare_wasm_instance(
-    global_context: &RequestGlobalContext,
     ctx: WasiCtx,
-    wasm_module_source: &WasmModuleSource,
+    wasm_module: &CompiledWasmModule,
     link_options: WasmLinkOptions,
 ) -> Result<(Store<WasiCtx>, Instance), Error> {
-    debug!("Creating store, engine, and linker.");
-    let (mut store, engine) = new_store_and_engine(&global_context.cache_config_path, ctx)?;
+    debug!("Cloning module object");
+    let (module, engine) = match wasm_module { CompiledWasmModule::Object(m, e) => (m.clone(), e.clone()) };
+    let mut store = new_store(ctx, &engine)?;
+
+    debug!("Configuring linker");
     let mut linker = Linker::new(&engine);
     wasmtime_wasi::add_to_linker(&mut linker, |cx| cx)?;
-
     link_options.apply_to(&mut linker)?;
 
-    debug!("loading module from store");
-    let module = wasm_module_source.load_module(&store)?;
     debug!("instantiating module in linker");
     let instance = linker.instantiate(&mut store, &module)?;
     Ok((store, instance))