make create_transform common & reorg trait imp to imp

xxchan · xxchan · Feb 26, 2023 · Feb 26, 2023 · Apr 9, 2023 · Feb 26, 2023
commit abd4fb61fb6df8e9858059a75569aaeb757dfd7e
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/fluvio-smartengine/src/engine/common.rs b/crates/fluvio-smartengine/src/engine/common.rs
@@ -8,9 +8,13 @@ use fluvio_smartmodule::dataplane::smartmodule::{
     SmartModuleExtraParams,
 };
 
-pub trait WasmInstance {
+use crate::SmartModuleInitialData;
+
+use super::error::EngineError;
+
+pub(crate) trait WasmInstance {
     type Context;
-    type Func: WasmFn<Context = Self::Context>;
+    type Func: WasmFn<Context = Self::Context> + Send + Sync + 'static;
 
     fn params(&self) -> SmartModuleExtraParams;
 
@@ -26,12 +30,12 @@ pub trait WasmInstance {
 
 /// All smartmodule wasm functions have the same ABI:
 /// `(ptr: *mut u8, len: usize, version: i16) -> i32`, which is `(param i32 i32 i32) (result i32)` in wasm.
-pub trait WasmFn {
+pub(crate) trait WasmFn {
     type Context;
     fn call(&self, ptr: i32, len: i32, version: i32, ctx: &mut Self::Context) -> Result<i32>;
 }
 
-pub trait SmartModuleTransform<I: WasmInstance>: Send + Sync {
+pub(crate) trait SmartModuleTransform<I: WasmInstance>: Send + Sync {
     /// transform records
     fn process(
         &mut self,
@@ -57,13 +61,13 @@ impl<T: SmartModuleTransform<I> + Any, I: WasmInstance> DowncastableTransform<I>
     }
 }
 
-pub struct SimpleTransformImpl<F: WasmFn + Send + Sync> {
+pub(crate) struct SimpleTransform<F: WasmFn + Send + Sync> {
     name: String,
     func: F,
 }
 
 impl<I: WasmInstance, F: WasmFn<Context = I::Context> + Send + Sync> SmartModuleTransform<I>
-    for SimpleTransformImpl<F>
+    for SimpleTransform<F>
 {
     fn process(
         &mut self,
@@ -89,15 +93,15 @@ impl<I: WasmInstance, F: WasmFn<Context = I::Context> + Send + Sync> SmartModule
     }
 }
 
-impl<F: WasmFn + Send + Sync> SimpleTransformImpl<F> {
-    pub(crate) fn try_instantiate<I>(
+impl<F: WasmFn + Send + Sync> SimpleTransform<F> {
+    pub(crate) fn try_instantiate<I, C>(
         name: &str,
         instance: &mut I,
-        ctx: &mut <I as WasmInstance>::Context,
+        ctx: &mut C,
     ) -> Result<Option<Self>>
     where
-        I: WasmInstance<Func = F>,
-        F: WasmFn<Context = I::Context>,
+        I: WasmInstance<Func = F, Context = C>,
+        F: WasmFn<Context = C>,
     {
         let func = instance
             .get_fn(name, ctx)?
@@ -109,6 +113,45 @@ impl<F: WasmFn + Send + Sync> SimpleTransformImpl<F> {
     }
 }
 
+pub(crate) const FILTER_FN_NAME: &str = "filter";
+pub(crate) const MAP_FN_NAME: &str = "map";
+pub(crate) const FILTER_MAP_FN_NAME: &str = "filter_map";
+pub(crate) const ARRAY_MAP_FN_NAME: &str = "array_map";
+
+pub(crate) fn create_transform<I, C>(
+    instance: &mut I,
+    ctx: &mut C,
+    _initial_data: SmartModuleInitialData,
+) -> Result<Box<dyn DowncastableTransform<I>>>
+where
+    I: WasmInstance<Context = C>,
+{
+    if let Some(tr) = SimpleTransform::try_instantiate(FILTER_FN_NAME, instance, ctx)?
+        .map(|transform| Box::new(transform) as Box<dyn DowncastableTransform<I>>)
+    {
+        Ok(tr)
+    } else if let Some(tr) = SimpleTransform::try_instantiate(MAP_FN_NAME, instance, ctx)?
+        .map(|transform| Box::new(transform) as Box<dyn DowncastableTransform<I>>)
+    {
+        Ok(tr)
+    } else if let Some(tr) = SimpleTransform::try_instantiate(FILTER_MAP_FN_NAME, instance, ctx)?
+        .map(|transform| Box::new(transform) as Box<dyn DowncastableTransform<I>>)
+    {
+        Ok(tr)
+    } else if let Some(tr) = SimpleTransform::try_instantiate(ARRAY_MAP_FN_NAME, instance, ctx)?
+        .map(|transform| Box::new(transform) as Box<dyn DowncastableTransform<I>>)
+    {
+        Ok(tr)
+        // TODO: AGGREGATE
+        // } else if let Some(tr) = SmartModuleAggregate::try_instantiate(ctx, initial_data, store)?
+        //     .map(|transform| Box::new(transform) as Box<dyn DowncastableTransform<I>>)
+        // {
+        //     Ok(tr)
+    } else {
+        Err(EngineError::UnknownSmartModule.into())
+    }
+}
+
 pub(crate) const INIT_FN_NAME: &str = "init";
 
 pub(crate) struct SmartModuleInit<F: WasmFn>(F);
@@ -220,87 +263,3 @@ impl<I: WasmInstance<Func = F>, F: WasmFn + Send + Sync> SmartModuleInstance<I,
         }
     }
 }
-
-mod wasmtime {
-    use anyhow::Result;
-    use fluvio_protocol::{Encoder, Decoder};
-    use tracing::debug;
-
-    use std::sync::Arc;
-
-    use fluvio_smartmodule::dataplane::smartmodule::SmartModuleExtraParams;
-    use wasmtime::{Instance, Store};
-
-    use crate::engine::wasmtime::instance::RecordsCallBack;
-
-    pub struct WasmTimeInstance {
-        instance: Instance,
-        records_cb: Arc<RecordsCallBack>,
-        params: SmartModuleExtraParams,
-        version: i16,
-    }
-    pub struct WasmTimeContext {
-        store: Store<()>,
-    }
-
-    pub type WasmTimeFn = wasmtime::TypedFunc<(i32, i32, i32), i32>;
-
-    impl super::WasmInstance for WasmTimeInstance {
-        type Context = WasmTimeContext;
-
-        type Func = WasmTimeFn;
-
-        fn get_fn(&self, name: &str, ctx: &mut Self::Context) -> Result<Option<Self::Func>> {
-            match self.instance.get_func(&mut ctx.store, name) {
-                Some(func) => {
-                    // check type signature
-                    func.typed(&mut ctx.store)
-                        .or_else(|_| func.typed(&ctx.store))
-                        .map(|f| Some(f))
-                }
-                None => Ok(None),
-            }
-        }
-
-        fn write_input<E: Encoder>(
-            &mut self,
-            input: &E,
-            ctx: &mut Self::Context,
-        ) -> anyhow::Result<(i32, i32, i32)> {
-            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 = crate::engine::wasmtime::memory::copy_memory_to_instance(
-                &mut ctx.store,
-                &self.instance,
-                &input_data,
-            )?;
-            let length = input_data.len();
-            Ok((array_ptr as i32, length as i32, self.version as i32))
-        }
-
-        fn read_output<D: Decoder + Default>(&mut self, ctx: &mut Self::Context) -> Result<D> {
-            let bytes = self
-                .records_cb
-                .get()
-                .and_then(|m| m.copy_memory_from(&ctx.store).ok())
-                .unwrap_or_default();
-            let mut output = D::default();
-            output.decode(&mut std::io::Cursor::new(bytes), self.version)?;
-            Ok(output)
-        }
-    }
-
-    impl super::WasmFn for WasmTimeFn {
-        type Context = WasmTimeContext;
-
-        fn call(&self, ptr: i32, len: i32, version: i32, ctx: &mut Self::Context) -> Result<i32> {
-            WasmTimeFn::call(self, &mut ctx.store, (ptr, len, version))
-        }
-    }
-}
diff --git a/crates/fluvio-smartengine/src/engine/wasmedge/imp.rs b/crates/fluvio-smartengine/src/engine/wasmedge/imp.rs
@@ -0,0 +1,153 @@
+use tracing::debug;
+use wasmedge_sdk::error::HostFuncError;
+use wasmedge_sdk::types::Val;
+use wasmedge_sdk::{
+    Executor, Func, Module, Store, CallingFrame, WasmValue, Caller, ImportObjectBuilder,
+};
+
+use crate::engine::wasmedge::instance::{RecordsCallBack, RecordsMemory};
+
+use crate::engine::common::{WasmFn, WasmInstance};
+use crate::engine::error::EngineError;
+
+use anyhow::Result;
+use fluvio_smartmodule::dataplane::smartmodule::{SmartModuleExtraParams};
+use std::sync::Arc;
+
+pub(crate) struct WasmedgeInstance {
+    instance: wasmedge_sdk::Instance,
+    records_cb: Arc<super::instance::RecordsCallBack>,
+    params: SmartModuleExtraParams,
+    version: i16,
+}
+
+pub(crate) struct WasmedgeContext {
+    pub engine: Executor,
+}
+
+pub type WasmedgeFn = Func;
+
+impl WasmInstance for WasmedgeInstance {
+    type Context = WasmedgeContext;
+
+    type Func = WasmedgeFn;
+
+    fn get_fn(&self, name: &str, _ctx: &mut Self::Context) -> Result<Option<Self::Func>> {
+        match self.instance.func(name) {
+            // check type signature
+            Some(func) => Ok(Some(func)),
+            None => Ok(None),
+        }
+    }
+
+    fn write_input<E: fluvio_protocol::Encoder>(
+        &mut self,
+        input: &E,
+        ctx: &mut Self::Context,
+    ) -> Result<(i32, i32, i32)> {
+        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 =
+            super::memory::copy_memory_to_instance(&ctx.engine, &self.instance, &input_data)?;
+        let length = input_data.len();
+        Ok((array_ptr, length as i32, self.version as i32))
+    }
+
+    fn read_output<D: fluvio_protocol::Decoder + Default>(
+        &mut self,
+        _ctx: &mut Self::Context,
+    ) -> 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)
+    }
+
+    fn params(&self) -> SmartModuleExtraParams {
+        self.params.clone()
+    }
+}
+
+impl WasmFn for WasmedgeFn {
+    type Context = WasmedgeContext;
+
+    fn call(&self, ptr: i32, len: i32, version: i32, ctx: &mut Self::Context) -> Result<i32> {
+        let res = self.call(
+            &ctx.engine,
+            vec![
+                Val::I32(ptr).into(),
+                Val::I32(len).into(),
+                Val::I32(version).into(),
+            ],
+        )?;
+        Ok(res[0].to_i32())
+    }
+}
+
+impl WasmedgeInstance {
+    /// 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,
+        })
+    }
+}
diff --git a/crates/fluvio-smartengine/src/engine/wasmedge/init.rs b/crates/fluvio-smartengine/src/engine/wasmedge/init.rs
@@ -0,0 +1,11 @@
+use std::convert::TryFrom;
+use std::fmt::Debug;
+
+use anyhow::{Result, Ok};
+use fluvio_smartmodule::dataplane::smartmodule::{
+    SmartModuleInitInput, SmartModuleInitOutput, SmartModuleInitErrorStatus,
+};
+
+use crate::engine::common::{WasmFn, WasmInstance};
+
+use super::instance::SmartModuleInstanceContext;