make safe context lifetime

Author: waralexrom

packages/cubejs-backend-native/src/node_export.rs

@@ -23,7 +23,7 @@ use crate::stream::OnDrainHandler;
 use crate::tokio_runtime_node;
 use crate::transport::NodeBridgeTransport;
 use crate::utils::batch_to_rows;
-use cubenativeutils::wrappers::neon::context::ContextHolder;
+use cubenativeutils::wrappers::neon::context::neon_run_with_guarded_lifetime;
 use cubenativeutils::wrappers::neon::inner_types::NeonInnerTypes;
 use cubenativeutils::wrappers::neon::object::NeonObject;
 use cubenativeutils::wrappers::object_handle::NativeObjectHandle;
@@ -463,38 +463,29 @@ pub fn setup_logger(mut cx: FunctionContext) -> JsResult<JsUndefined> {
 //============ sql planner ===================
 
 fn build_sql_and_params(cx: FunctionContext) -> JsResult<JsValue> {
-    //IMPORTANT It seems to be safe here, because context lifetime is bound to function, but this
-    //context should be used only inside function
-    let mut cx = extend_function_context_lifetime(cx);
-    let options = cx.argument::<JsValue>(0)?;
-    let options_tmp = options.clone();
-
-    let neon_context_holder = ContextHolder::new(cx);
-
-    let options = NativeObjectHandle::<NeonInnerTypes<'static, FunctionContext<'static>>>::new(
-        NeonObject::new(neon_context_holder.clone(), options),
-    );
-
-    let context_holder =
-        NativeContextHolder::<NeonInnerTypes<'static, FunctionContext<'static>>>::new(
+    neon_run_with_guarded_lifetime(cx, |neon_context_holder| {
+        let options =
+            NativeObjectHandle::<NeonInnerTypes<FunctionContext<'static>>>::new(NeonObject::new(
+                neon_context_holder.clone(),
+                neon_context_holder
+                    .with_context(|cx| cx.argument::<JsValue>(0))
+                    .unwrap()?,
+            ));
+
+        let context_holder = NativeContextHolder::<NeonInnerTypes<FunctionContext<'static>>>::new(
             neon_context_holder,
         );
 
-    let base_query_options = Rc::new(NativeBaseQueryOptions::from_native(options).unwrap());
+        let base_query_options = Rc::new(NativeBaseQueryOptions::from_native(options).unwrap());
 
-    let base_query = BaseQuery::try_new(context_holder.clone(), base_query_options).unwrap();
+        let base_query = BaseQuery::try_new(context_holder.clone(), base_query_options).unwrap();
 
-    //arg_clrep.into_js(&mut cx)
-    let res = base_query.build_sql_and_params();
-
-    let result: NeonObject<'static, FunctionContext<'static>> = res.into_object();
-    let result = result.into_object();
-
-    Ok(result)
-}
+        let res = base_query.build_sql_and_params();
 
-fn extend_function_context_lifetime<'a>(cx: FunctionContext<'a>) -> FunctionContext<'static> {
-    unsafe { std::mem::transmute::<FunctionContext<'a>, FunctionContext<'static>>(cx) }
+        let result: NeonObject<FunctionContext<'static>> = res.into_object();
+        let result = result.into_object();
+        Ok(result)
+    })
 }
 
 fn debug_js_to_clrepr_to_js(mut cx: FunctionContext) -> JsResult<JsValue> {

rust/cubenativeutils/src/wrappers/context.rs

@@ -1,16 +1,17 @@
 use super::inner_types::InnerTypes;
 use super::object::NativeBox;
 use super::object_handle::NativeObjectHandle;
+use cubesql::CubeError;
 
 pub trait NativeContext<IT: InnerTypes>: Clone {
-    fn boolean(&self, v: bool) -> IT::Boolean;
-    fn string(&self, v: String) -> IT::String;
-    fn number(&self, v: f64) -> IT::Number;
-    fn undefined(&self) -> NativeObjectHandle<IT>;
-    fn empty_array(&self) -> IT::Array;
-    fn empty_struct(&self) -> IT::Struct;
+    fn boolean(&self, v: bool) -> Result<IT::Boolean, CubeError>;
+    fn string(&self, v: String) -> Result<IT::String, CubeError>;
+    fn number(&self, v: f64) -> Result<IT::Number, CubeError>;
+    fn undefined(&self) -> Result<NativeObjectHandle<IT>, CubeError>;
+    fn empty_array(&self) -> Result<IT::Array, CubeError>;
+    fn empty_struct(&self) -> Result<IT::Struct, CubeError>;
     //fn boxed<T: 'static>(&self, value: T) -> impl NativeBox<IT, T>;
-    fn to_string_fn(&self, result: String) -> IT::Function;
+    fn to_string_fn(&self, result: String) -> Result<IT::Function, CubeError>;
 }
 
 #[derive(Clone)]
@@ -25,26 +26,26 @@ impl<IT: InnerTypes> NativeContextHolder<IT> {
     pub fn context(&self) -> &impl NativeContext<IT> {
         &self.context
     }
-    pub fn boolean(&self, v: bool) -> IT::Boolean {
+    pub fn boolean(&self, v: bool) -> Result<IT::Boolean, CubeError> {
         self.context.boolean(v)
     }
-    pub fn string(&self, v: String) -> IT::String {
+    pub fn string(&self, v: String) -> Result<IT::String, CubeError> {
         self.context.string(v)
     }
-    pub fn number(&self, v: f64) -> IT::Number {
+    pub fn number(&self, v: f64) -> Result<IT::Number, CubeError> {
         self.context.number(v)
     }
-    pub fn undefined(&self) -> NativeObjectHandle<IT> {
+    pub fn undefined(&self) -> Result<NativeObjectHandle<IT>, CubeError> {
         self.context.undefined()
     }
-    pub fn empty_array(&self) -> IT::Array {
+    pub fn empty_array(&self) -> Result<IT::Array, CubeError> {
         self.context.empty_array()
     }
-    pub fn empty_struct(&self) -> IT::Struct {
+    pub fn empty_struct(&self) -> Result<IT::Struct, CubeError> {
         self.context.empty_struct()
     }
     #[allow(dead_code)]
-    pub fn to_string_fn(&self, result: String) -> IT::Function {
+    pub fn to_string_fn(&self, result: String) -> Result<IT::Function, CubeError> {
         self.context.to_string_fn(result)
     }
 }

rust/cubenativeutils/src/wrappers/neon/context.rs

@@ -10,20 +10,29 @@ use crate::wrappers::object::NativeObject;
 use crate::wrappers::object_handle::NativeObjectHandle;
 use cubesql::CubeError;
 use neon::prelude::*;
-use std::cell::{RefCell, RefMut};
-use std::marker::PhantomData;
+use std::cell::RefCell;
+use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
 use std::rc::{Rc, Weak};
-pub struct ContextWrapper<'cx, C: Context<'cx>> {
+
+pub trait NoenContextLifetimeExpand<'cx> {
+    type ExpandedResult: Context<'static>;
+    fn expand_lifetime(self) -> Self::ExpandedResult;
+}
+
+impl<'cx> NoenContextLifetimeExpand<'cx> for FunctionContext<'cx> {
+    type ExpandedResult = FunctionContext<'static>;
+    fn expand_lifetime(self) -> Self::ExpandedResult {
+        unsafe { std::mem::transmute::<FunctionContext<'cx>, FunctionContext<'static>>(self) }
+    }
+}
+
+pub struct ContextWrapper<C: Context<'static>> {
     cx: C,
-    lifetime: PhantomData<&'cx ()>,
 }
 
-impl<'cx, C: Context<'cx>> ContextWrapper<'cx, C> {
+impl<C: Context<'static>> ContextWrapper<C> {
     pub fn new(cx: C) -> Rc<RefCell<Self>> {
-        Rc::new(RefCell::new(Self {
-            cx,
-            lifetime: Default::default(),
-        }))
+        Rc::new(RefCell::new(Self { cx }))
     }
 
     pub fn with_context<T, F>(&mut self, f: F) -> T
@@ -38,125 +47,105 @@ impl<'cx, C: Context<'cx>> ContextWrapper<'cx, C> {
     }
 }
 
-pub struct ContextHolder<'cx, C: Context<'cx>> {
-    context: Rc<RefCell<ContextWrapper<'cx, C>>>,
+pub fn neon_run_with_guarded_lifetime<'cx, C, T, F>(cx: C, func: F) -> T
+where
+    C: Context<'cx> + NoenContextLifetimeExpand<'cx>,
+    F: FnOnce(ContextHolder<C::ExpandedResult>) -> T,
+{
+    let context = ContextWrapper::new(cx.expand_lifetime());
+    let context_holder = ContextHolder::new(Rc::downgrade(&context));
+    let res = catch_unwind(AssertUnwindSafe(|| func(context_holder)));
+    match Rc::try_unwrap(context) {
+        Ok(_) => {}
+        Err(_) => panic!("Guarded context have more then one reference"),
+    };
+
+    match res {
+        Ok(res) => res,
+        Err(e) => resume_unwind(e),
+    }
 }
 
-impl<'cx, C: Context<'cx> + 'cx> ContextHolder<'cx, C> {
-    pub fn new(cx: C) -> Self {
-        Self {
-            context: ContextWrapper::new(cx),
-        }
-    }
+pub struct ContextHolder<C: Context<'static>> {
+    context: Weak<RefCell<ContextWrapper<C>>>,
+}
 
-    pub fn borrow_mut(&self) -> RefMut<ContextWrapper<'cx, C>> {
-        self.context.borrow_mut()
+impl<C: Context<'static>> ContextHolder<C> {
+    fn new(context: Weak<RefCell<ContextWrapper<C>>>) -> Self {
+        Self { context }
     }
 
-    pub fn with_context<T, F>(&self, f: F) -> T
+    pub fn with_context<T, F>(&self, f: F) -> Result<T, CubeError>
     where
         F: FnOnce(&mut C) -> T,
     {
-        let mut context = self.context.borrow_mut();
-        context.with_context(f)
-    }
-
-    pub fn weak(&self) -> WeakContextHolder<'cx, C> {
-        WeakContextHolder {
-            context: Rc::downgrade(&self.context),
+        if let Some(context) = self.context.upgrade() {
+            let mut cx = context.borrow_mut();
+            let res = cx.with_context(f);
+            Ok(res)
+        } else {
+            Err(CubeError::internal(format!(
+                "Call to neon context outside of its lifetime"
+            )))
         }
     }
 }
 
-impl<'cx, C: Context<'cx> + 'cx> NativeContext<NeonInnerTypes<'cx, C>> for ContextHolder<'cx, C> {
-    fn boolean(&self, v: bool) -> NeonBoolean<'cx, C> {
-        let obj = NeonObject::new(self.clone(), self.with_context(|cx| cx.boolean(v).upcast()));
-        obj.into_boolean().unwrap()
+impl<C: Context<'static> + 'static> NativeContext<NeonInnerTypes<C>> for ContextHolder<C> {
+    fn boolean(&self, v: bool) -> Result<NeonBoolean<C>, CubeError> {
+        let obj = NeonObject::new(
+            self.clone(),
+            self.with_context(|cx| cx.boolean(v).upcast())?,
+        );
+        obj.into_boolean()
     }
 
-    fn string(&self, v: String) -> NeonString<'cx, C> {
-        let obj = NeonObject::new(self.clone(), self.with_context(|cx| cx.string(v).upcast()));
-        obj.into_string().unwrap()
+    fn string(&self, v: String) -> Result<NeonString<C>, CubeError> {
+        let obj = NeonObject::new(self.clone(), self.with_context(|cx| cx.string(v).upcast())?);
+        obj.into_string()
     }
 
-    fn number(&self, v: f64) -> NeonNumber<'cx, C> {
-        let obj = NeonObject::new(self.clone(), self.with_context(|cx| cx.number(v).upcast()));
-        obj.into_number().unwrap()
+    fn number(&self, v: f64) -> Result<NeonNumber<C>, CubeError> {
+        let obj = NeonObject::new(self.clone(), self.with_context(|cx| cx.number(v).upcast())?);
+        obj.into_number()
     }
 
-    fn undefined(&self) -> NativeObjectHandle<NeonInnerTypes<'cx, C>> {
-        NativeObjectHandle::new(NeonObject::new(
+    fn undefined(&self) -> Result<NativeObjectHandle<NeonInnerTypes<C>>, CubeError> {
+        Ok(NativeObjectHandle::new(NeonObject::new(
             self.clone(),
-            self.with_context(|cx| cx.undefined().upcast()),
-        ))
+            self.with_context(|cx| cx.undefined().upcast())?,
+        )))
     }
 
-    fn empty_array(&self) -> NeonArray<'cx, C> {
+    fn empty_array(&self) -> Result<NeonArray<C>, CubeError> {
         let obj = NeonObject::new(
             self.clone(),
-            self.with_context(|cx| cx.empty_array().upcast()),
+            self.with_context(|cx| cx.empty_array().upcast())?,
         );
-        obj.into_array().unwrap()
+        obj.into_array()
     }
 
-    fn empty_struct(&self) -> NeonStruct<'cx, C> {
+    fn empty_struct(&self) -> Result<NeonStruct<C>, CubeError> {
         let obj = NeonObject::new(
             self.clone(),
-            self.with_context(|cx| cx.empty_object().upcast()),
+            self.with_context(|cx| cx.empty_object().upcast())?,
         );
-        obj.into_struct().unwrap()
+        obj.into_struct()
     }
-    fn to_string_fn(&self, result: String) -> NeonFunction<'cx, C> {
+    fn to_string_fn(&self, result: String) -> Result<NeonFunction<C>, CubeError> {
         let obj = NeonObject::new(
             self.clone(),
             self.with_context(|cx| {
                 JsFunction::new(cx, move |mut c| Ok(c.string(result.clone())))
                     .unwrap()
                     .upcast()
-            }),
+            })?,
         );
-        obj.into_function().unwrap()
-    }
-}
-
-impl<'cx, C: Context<'cx>> Clone for ContextHolder<'cx, C> {
-    fn clone(&self) -> Self {
-        Self {
-            context: self.context.clone(),
-        }
-    }
-}
-
-pub struct WeakContextHolder<'cx, C: Context<'cx>> {
-    context: Weak<RefCell<ContextWrapper<'cx, C>>>,
-}
-
-impl<'cx, C: Context<'cx>> WeakContextHolder<'cx, C> {
-    pub fn try_upgrade<'a>(&'a self) -> Result<ContextHolder<'a, C>, CubeError>
-    where
-        'a: 'cx,
-    {
-        if let Some(context) = self.context.upgrade() {
-            Ok(ContextHolder { context })
-        } else {
-            Err(CubeError::internal(format!("Neon context is not alive")))
-        }
-    }
-    pub fn with_context<T, F>(&self, f: F) -> Result<T, CubeError>
-    where
-        F: FnOnce(&mut C) -> T,
-    {
-        if let Some(context) = self.context.upgrade() {
-            let mut cx = context.borrow_mut();
-            let res = cx.with_context(f);
-            Ok(res)
-        } else {
-            Err(CubeError::internal(format!("Neon context is not alive")))
-        }
+        obj.into_function()
     }
 }
 
-impl<'cx, C: Context<'cx>> Clone for WeakContextHolder<'cx, C> {
+impl<C: Context<'static>> Clone for ContextHolder<C> {
     fn clone(&self) -> Self {
         Self {
             context: self.context.clone(),

rust/cubenativeutils/src/wrappers/neon/inner_types.rs

@@ -8,25 +8,25 @@ use crate::wrappers::inner_types::InnerTypes;
 use neon::prelude::*;
 use std::marker::PhantomData;
 
-pub struct NeonInnerTypes<'cx: 'static, C: Context<'cx>> {
-    lifetime: PhantomData<&'cx ContextHolder<'cx, C>>,
+pub struct NeonInnerTypes<C: Context<'static>> {
+    marker: PhantomData<ContextHolder<C>>,
 }
 
-impl<'cx, C: Context<'cx>> Clone for NeonInnerTypes<'cx, C> {
+impl<C: Context<'static>> Clone for NeonInnerTypes<C> {
     fn clone(&self) -> Self {
         Self {
-            lifetime: Default::default(),
+            marker: Default::default(),
         }
     }
 }
 
-impl<'cx: 'static, C: Context<'cx>> InnerTypes for NeonInnerTypes<'cx, C> {
-    type Object = NeonObject<'cx, C>;
-    type Context = ContextHolder<'cx, C>;
-    type Array = NeonArray<'cx, C>;
-    type Struct = NeonStruct<'cx, C>;
-    type String = NeonString<'cx, C>;
-    type Boolean = NeonBoolean<'cx, C>;
-    type Function = NeonFunction<'cx, C>;
-    type Number = NeonNumber<'cx, C>;
+impl<C: Context<'static> + 'static> InnerTypes for NeonInnerTypes<C> {
+    type Object = NeonObject<C>;
+    type Context = ContextHolder<C>;
+    type Array = NeonArray<C>;
+    type Struct = NeonStruct<C>;
+    type String = NeonString<C>;
+    type Boolean = NeonBoolean<C>;
+    type Function = NeonFunction<C>;
+    type Number = NeonNumber<C>;
 }