support init for wasmedge/common

xxchan · xxchan · Feb 26, 2023 · Feb 26, 2023 · Apr 9, 2023 · Feb 26, 2023
commit f17b8ef6322129e965d46da6abf1e275982bba8c
diff --git a/crates/fluvio-smartengine/Cargo.toml b/crates/fluvio-smartengine/Cargo.toml
@@ -12,11 +12,12 @@ description = "The official Fluvio SmartEngine"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-engine = ["wasmtime-engine"]
+engine = ["wasmtime-engine", "wasmedge-engine"]
 wasi = ["wasmtime-wasi", "wasmtime-engine", "engine"]
 transformation = ["serde_json", "serde_yaml"]
 default = ["engine"]
 wasmtime-engine = ["wasmtime"]
+wasmedge-engine = ["wasmedge-sdk"]
 
 [dependencies]
 tracing ={ workspace = true }

diff --git a/crates/fluvio-smartengine/src/engine/common.rs b/crates/fluvio-smartengine/src/engine/common.rs
@@ -4,6 +4,7 @@ use anyhow::Result;
 use fluvio_protocol::{Encoder, Decoder};
 use fluvio_smartmodule::dataplane::smartmodule::{
     SmartModuleInput, SmartModuleOutput, SmartModuleTransformErrorStatus,
+    SmartModuleInitErrorStatus, SmartModuleInitOutput, SmartModuleInitInput,
 };
 
 pub trait WasmInstance {
@@ -105,6 +106,64 @@ impl<F: WasmFn + Send + Sync> SimpleTransformImpl<F> {
     }
 }
 
+pub(crate) const INIT_FN_NAME: &str = "init";
+
+pub(crate) struct SmartModuleInit<F: WasmFn>(F);
+
+impl<F: WasmFn> std::fmt::Debug for SmartModuleInit<F> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "InitFn")
+    }
+}
+
+impl<F: WasmFn + Send + Sync> SmartModuleInit<F> {
+    pub(crate) fn try_instantiate<I>(
+        instance: &mut I,
+        ctx: &mut <I as WasmInstance>::Context,
+    ) -> Result<Option<Self>>
+    where
+        I: WasmInstance<Func = F>,
+        F: WasmFn<Context = I::Context>,
+    {
+        match instance.get_fn(INIT_FN_NAME, ctx)? {
+            Some(func) => Ok(Some(Self(func))),
+            None => Ok(None),
+        }
+    }
+}
+
+impl<F: WasmFn + Send + Sync> SmartModuleInit<F> {
+    /// initialize SmartModule
+    pub(crate) fn initialize<I>(
+        &mut self,
+        input: SmartModuleInitInput,
+        instance: &mut I,
+        ctx: &mut I::Context,
+    ) -> Result<()>
+    where
+        I: WasmInstance,
+        F: WasmFn<Context = I::Context>,
+    {
+        let (ptr, len, version) = instance.write_input(&input, ctx)?;
+        let init_output = self.0.call(ptr, len, version, ctx)?;
+
+        if init_output < 0 {
+            let internal_error = SmartModuleInitErrorStatus::try_from(init_output)
+                .unwrap_or(SmartModuleInitErrorStatus::UnknownError);
+
+            match internal_error {
+                SmartModuleInitErrorStatus::InitError => {
+                    let output: SmartModuleInitOutput = instance.read_output(ctx)?;
+                    Err(output.error.into())
+                }
+                _ => Err(internal_error.into()),
+            }
+        } else {
+            Ok(())
+        }
+    }
+}
+
 mod wasmtime {
     use anyhow::Result;
     use fluvio_protocol::{Encoder, Decoder};

diff --git a/crates/fluvio-smartengine/src/engine/wasmedge/init.rs b/crates/fluvio-smartengine/src/engine/wasmedge/init.rs
diff --git a/crates/fluvio-smartengine/src/engine/wasmedge/instance.rs b/crates/fluvio-smartengine/src/engine/wasmedge/instance.rs
@@ -8,22 +8,27 @@ use wasmedge_sdk::{
 };
 
 use super::{WasmedgeInstance, WasmedgeContext};
-use super::init::SmartModuleInit;
+use crate::engine::common::SmartModuleInit;
 use crate::engine::common::DowncastableTransform;
 use crate::engine::error::EngineError;
 use crate::metrics::SmartModuleChainMetrics;
 use crate::engine::wasmedge::memory;
 use crate::engine::{config::*, WasmSlice};
 use anyhow::Result;
-use fluvio_smartmodule::dataplane::smartmodule::{SmartModuleInput, SmartModuleOutput};
+use fluvio_smartmodule::dataplane::smartmodule::{
+    SmartModuleInput, SmartModuleOutput, SmartModuleInitInput, SmartModuleExtraParams,
+};
 use std::any::Any;
 use std::fmt::{self, Debug};
 use std::sync::{Arc, Mutex};
+use super::WasmedgeFn;
+
+use super::Init;
 
 pub(crate) struct SmartModuleInstance {
     pub instance: WasmedgeInstance,
     pub transform: Box<dyn DowncastableTransform<WasmedgeInstance>>,
-    // init: Option<SmartModuleInit>,
+    pub init: Option<Init>,
 }
 
 impl SmartModuleInstance {
@@ -34,18 +39,18 @@ impl SmartModuleInstance {
     }
 
     #[cfg(test)]
-    pub(crate) fn get_init(&self) -> &Option<SmartModuleInit> {
-        &None
+    pub(crate) fn get_init(&self) -> &Option<Init> {
+        &self.init
     }
 
     pub(crate) fn new(
         instance: WasmedgeInstance,
-        // init: Option<SmartModuleInit>,
+        init: Option<Init>,
         transform: Box<dyn DowncastableTransform<WasmedgeInstance>>,
     ) -> Self {
         Self {
             instance,
-            // init,
+            init,
             transform,
         }
     }
@@ -57,108 +62,16 @@ impl SmartModuleInstance {
     ) -> Result<SmartModuleOutput> {
         self.transform.process(input, &mut self.instance, ctx)
     }
-}
-
-pub struct SmartModuleInstanceContext {
-    instance: Instance,
-    records_cb: Arc<RecordsCallBack>,
-    // params: SmartModuleExtraParams,
-    version: i16,
-}
-
-impl SmartModuleInstanceContext {
-    /// get wasm function from instance
-    pub(crate) fn get_wasm_func(&self, name: &str) -> Option<Func> {
-        self.instance.func(name)
-    }
-
-    /// instantiate new module instance that contain context
-    pub(crate) fn instantiate(
-        store: &mut Store,
-        executor: &mut Executor,
-        module: Module,
-        // params: SmartModuleExtraParams,
-        version: i16,
-    ) -> Result<Self, EngineError> {
-        debug!("creating WasmModuleInstance");
-        let cb = Arc::new(RecordsCallBack::new());
-        let records_cb = cb.clone();
-
-        // See crates/fluvio-smartmodule-derive/src/generator/transform.rs for copy_records
-        let copy_records_fn = move |caller: CallingFrame,
-                                    inputs: Vec<WasmValue>|
-              -> Result<Vec<WasmValue>, HostFuncError> {
-            assert_eq!(inputs.len(), 2);
-            let ptr = inputs[0].to_i32() as u32;
-            let len = inputs[1].to_i32() as u32;
-
-            debug!(len, "callback from wasm filter");
-            let caller = Caller::new(caller);
-            let memory = caller.memory(0).unwrap();
-
-            let records = RecordsMemory { ptr, len, memory };
-            cb.set(records);
-            Ok(vec![])
-        };
 
-        let import = ImportObjectBuilder::new()
-            .with_func::<(i32, i32), ()>("copy_records", copy_records_fn)
-            .map_err(|e| EngineError::Instantiate(e.into()))?
-            .build("env")
-            .map_err(|e| EngineError::Instantiate(e.into()))?;
-
-        debug!("instantiating WASMtime");
-        store
-            .register_import_module(executor, &import)
-            .map_err(|e| EngineError::Instantiate(e.into()))?;
-        let instance = store
-            .register_active_module(executor, &module)
-            .map_err(|e| EngineError::Instantiate(e.into()))?;
-
-        // This is a hack to avoid them being dropped
-        // FIXME: manage their lifetimes
-        std::mem::forget(import);
-        std::mem::forget(module);
-
-        Ok(Self {
-            instance,
-            records_cb,
-            // params,
-            version,
-        })
-    }
-
-    pub(crate) fn write_input<E: Encoder>(
-        &mut self,
-        input: &E,
-        engine: &impl Engine,
-    ) -> Result<Vec<WasmValue>> {
-        self.records_cb.clear();
-        let mut input_data = Vec::new();
-        input.encode(&mut input_data, self.version)?;
-        debug!(
-            len = input_data.len(),
-            version = self.version,
-            "input encoded"
-        );
-        let array_ptr = memory::copy_memory_to_instance(engine, &self.instance, &input_data)?;
-        let length = input_data.len();
-        Ok(vec![
-            Val::I32(array_ptr as i32).into(),
-            Val::I32(length as i32).into(),
-            Val::I32(self.version as i32).into(),
-        ])
-    }
-
-    pub(crate) fn read_output<D: Decoder + Default>(&mut self) -> Result<D> {
-        let bytes = self
-            .records_cb
-            .get()
-            .and_then(|m| m.copy_memory_from().ok())
-            .unwrap_or_default();
-        let mut output = D::default();
-        output.decode(&mut std::io::Cursor::new(bytes), self.version)?;
-        Ok(output)
+    pub fn init(&mut self, ctx: &mut WasmedgeContext) -> Result<()> {
+        if let Some(init) = &mut self.init {
+            let input = SmartModuleInitInput {
+                params: self.instance.params.clone(),
+            };
+            init.initialize(input, &mut self.instance, ctx)
+        } else {
+            Ok(())
+        }
     }
 }
 

diff --git a/crates/fluvio-smartengine/src/engine/wasmedge/mod.rs b/crates/fluvio-smartengine/src/engine/wasmedge/mod.rs
@@ -1,8 +1,6 @@
 mod instance;
 mod transforms;
 use instance::*;
-mod init;
-use init::*;
 mod memory;
 use memory::*;
 
@@ -28,6 +26,10 @@ use std::sync::{Arc, Mutex};
 
 use self::transforms::create_transform;
 
+use super::common::SmartModuleInit;
+
+type Init = SmartModuleInit<WasmedgeFn>;
+
 pub struct WasmedgeInstance {
     instance: wasmedge_sdk::Instance,
     records_cb: Arc<RecordsCallBack>,
@@ -208,13 +210,15 @@ impl SmartModuleChainBuilder {
                 version,
             )?;
 
-            // let init = SmartModuleInit::try_instantiate(&ctx, &mut state)?;
+            let init = Init::try_instantiate(&mut instance, &mut ctx)?;
             let transform = create_transform(&mut instance, &mut ctx, config.initial_data)?;
-            // instance.init(&mut state)?;
-            instances.push(SmartModuleInstance {
+            let mut instance = SmartModuleInstance {
                 instance,
                 transform,
-            });
+                init,
+            };
+            instance.init(&mut ctx)?;
+            instances.push(instance);
         }
 
         Ok(SmartModuleChainInstance { ctx, instances })