feat(native): Jinja - async render (#7309)

Author: ovr

packages/cubejs-backend-native/Cargo.lock

@@ -15,7 +15,8 @@ crate-type = ["cdylib"]
 [dependencies]
 cubesql = { path = "../../rust/cubesql/cubesql" }
 cubeclient = { path = "../../rust/cubesql/cubeclient" }
-tokio = { version = "1.0", features = ["full", "rt"] }
+tokio = { version = "1", features = ["full", "rt"] }
+async-channel = { version = "2" }
 async-trait = "0.1.36"
 serde_derive = "1.0.115"
 serde = "1.0.115"

packages/cubejs-backend-native/Cargo.toml

@@ -376,9 +376,15 @@ export type PythonCtx = {
   variables: Record<string, any>
 };
 
+export type JinjaEngineOptions = {
+  debugInfo?: boolean,
+  filters: Record<string, Function>,
+  workers: number
+};
+
 export interface JinjaEngine {
     loadTemplate(templateName: string, templateContent: string): void;
-    renderTemplate(templateName: string, context: unknown, pythonContext: Record<string, any> | null): string;
+    renderTemplate(templateName: string, context: unknown, pythonContext: Record<string, any> | null): Promise<string>;
 }
 
 export class NativeInstance {
@@ -394,7 +400,7 @@ export class NativeInstance {
       return this.native;
     }
 
-    public newJinjaEngine(options: { debugInfo?: boolean, filters: Record<string, Function> }): JinjaEngine {
+    public newJinjaEngine(options: JinjaEngineOptions): JinjaEngine {
       return this.getNative().newJinjaEngine(options);
     }
 

packages/cubejs-backend-native/js/index.ts

@@ -1,4 +1,5 @@
 #![feature(async_closure)]
+#![feature(thread_id_value)]
 
 extern crate findshlibs;
 

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

@@ -1,18 +1,21 @@
 use crate::cross::*;
 use crate::template::mj_value::*;
-use crate::template::neon::NeonMiniJinjaContext;
+use crate::template::neon::*;
 use crate::utils::bind_method;
 
 use log::trace;
 use minijinja as mj;
 use neon::prelude::*;
 use std::cell::RefCell;
 
+use crate::template::workers::{JinjaEngineWorkerJob, JinjaEngineWorkerPool};
 #[cfg(feature = "python")]
 use pyo3::{exceptions::PyNotImplementedError, prelude::*, types::PyTuple, AsPyPointer};
 
 struct JinjaEngine {
     inner: mj::Environment<'static>,
+    workers_count: usize,
+    workers: Option<JinjaEngineWorkerPool>,
 }
 
 impl Finalize for JinjaEngine {}
@@ -123,14 +126,50 @@ impl JinjaEngine {
             }
         }
 
-        Ok(Self { inner: engine })
+        let workers_count = {
+            let workers_count_float = options
+                .get_value(cx, "workers")?
+                .downcast_or_throw::<JsNumber, _>(cx)?
+                .value(cx);
+
+            if workers_count_float < 1_f64 {
+                return cx.throw_error("Option workers must be a positive integer");
+            }
+
+            match workers_count_float.to_string().parse::<usize>() {
+                Ok(v) => v,
+                Err(err) => {
+                    return cx.throw_error(format!("Option workers must be a positive: {}", err))
+                }
+            }
+        };
+
+        Ok(Self {
+            inner: engine,
+            workers_count,
+            workers: None,
+        })
     }
 }
 
 type BoxedJinjaEngine = JsBox<RefCell<JinjaEngine>>;
 
 impl JinjaEngine {
-    fn render_template(mut cx: FunctionContext) -> JsResult<JsString> {
+    fn build_if_needed(&mut self, cx: &mut FunctionContext) -> &JinjaEngineWorkerPool {
+        if let Some(ref workers) = self.workers {
+            return workers;
+        }
+
+        self.workers = Some(JinjaEngineWorkerPool::new(
+            self.workers_count,
+            cx.channel(),
+            self.inner.clone(),
+        ));
+
+        self.workers.as_ref().unwrap()
+    }
+
+    fn render_template(mut cx: FunctionContext) -> JsResult<JsPromise> {
         #[cfg(build = "debug")]
         trace!("JinjaEngine.render_template");
 
@@ -142,16 +181,6 @@ impl JinjaEngine {
         let template_compile_context = CLRepr::from_js_ref(cx.argument::<JsValue>(1)?, &mut cx)?;
         let template_python_context = CLRepr::from_js_ref(cx.argument::<JsValue>(2)?, &mut cx)?;
 
-        let engine = &this.borrow().inner;
-        let template = match engine.get_template(&template_name.value(&mut cx)) {
-            Ok(t) => t,
-            Err(err) => {
-                trace!("jinja get template error: {:?}", err);
-
-                return cx.throw_from_mj_error(err);
-            }
-        };
-
         let mut to_jinja_ctx = CLReprObject::new();
         to_jinja_ctx.insert("COMPILE_CONTEXT".to_string(), template_compile_context);
 
@@ -163,15 +192,20 @@ impl JinjaEngine {
             }
         }
 
-        let compile_context = to_minijinja_value(CLRepr::Object(to_jinja_ctx));
-        match template.render(compile_context) {
-            Ok(r) => Ok(cx.string(r)),
-            Err(err) => {
-                trace!("jinja render template error: {:?}", err);
+        let (deferred, promise) = cx.promise();
 
-                cx.throw_from_mj_error(err)
-            }
-        }
+        let mut this = this.borrow_mut();
+        let pool = this.build_if_needed(&mut cx);
+
+        if let Err(err) = pool.render(JinjaEngineWorkerJob {
+            template_name: template_name.value(&mut cx),
+            ctx: to_minijinja_value(CLRepr::Object(to_jinja_ctx)),
+            deferred,
+        }) {
+            return cx.throw_error(format!("Unable to render jinja template: {}", err));
+        };
+
+        Ok(promise)
     }
 
     fn load_template(mut cx: FunctionContext) -> JsResult<JsUndefined> {
@@ -185,13 +219,18 @@ impl JinjaEngine {
         let template_name = cx.argument::<JsString>(0)?;
         let template_content = cx.argument::<JsString>(1)?;
 
-        if let Err(err) = this.borrow_mut().inner.add_template_owned(
+        let mut borrowed = this.borrow_mut();
+        if let Err(err) = borrowed.inner.add_template_owned(
             template_name.value(&mut cx),
             template_content.value(&mut cx),
         ) {
             trace!("jinja load error: {:?}", err);
-
             return cx.throw_from_mj_error(err);
+        };
+
+        if borrowed.workers.is_some() {
+            trace!("Restart jinja workers");
+            borrowed.workers = None;
         }
 
         Ok(cx.undefined())

packages/cubejs-backend-native/src/template/entry.rs

@@ -1,5 +1,6 @@
 mod entry;
 mod mj_value;
 mod neon;
+mod workers;
 
 pub use entry::template_register_module;

packages/cubejs-backend-native/src/template/mod.rs

@@ -0,0 +1,102 @@
+use crate::template::neon::*;
+use cubesql::CubeError;
+
+use log::trace;
+use minijinja as mj;
+use neon::prelude::*;
+use neon::types::Deferred;
+
+pub struct JinjaEngineWorkerJob {
+    pub(crate) template_name: String,
+    pub(crate) ctx: minijinja::value::Value,
+    pub(crate) deferred: Deferred,
+}
+
+struct JinjaEngineWorker {
+    _thread: std::thread::JoinHandle<()>,
+}
+
+impl JinjaEngineWorker {
+    fn new(
+        id: usize,
+        env: mj::Environment<'static>,
+        js_channel: neon::event::Channel,
+        receiver: async_channel::Receiver<JinjaEngineWorkerJob>,
+    ) -> Self {
+        let thread = std::thread::spawn(move || loop {
+            if let Ok(job) = receiver.recv_blocking() {
+                let template = match env.get_template(&job.template_name) {
+                    Ok(t) => t,
+                    Err(err) => {
+                        job.deferred.settle_with(
+                            &js_channel,
+                            move |mut cx| -> NeonResult<Handle<JsString>> {
+                                cx.throw_from_mj_error(err)
+                            },
+                        );
+
+                        continue;
+                    }
+                };
+
+                let result = template.render(job.ctx);
+                job.deferred.settle_with(
+                    &js_channel,
+                    move |mut cx| -> NeonResult<Handle<JsString>> {
+                        match result {
+                            Ok(r) => Ok(cx.string(r)),
+                            Err(err) => cx.throw_from_mj_error(err),
+                        }
+                    },
+                );
+            } else {
+                trace!(
+                    "Closing jinja thread, id: {}, threadId: {}",
+                    id,
+                    std::thread::current().id().as_u64()
+                );
+
+                return;
+            }
+        });
+
+        Self { _thread: thread }
+    }
+}
+
+pub struct JinjaEngineWorkerPool {
+    workers_rx: async_channel::Sender<JinjaEngineWorkerJob>,
+    _workers: Vec<JinjaEngineWorker>,
+}
+
+impl JinjaEngineWorkerPool {
+    pub fn new(
+        workers_count: usize,
+        js_channel: Channel,
+        jinja_engine: minijinja::Environment<'static>,
+    ) -> Self {
+        let (workers_rx, receiver) = async_channel::bounded::<JinjaEngineWorkerJob>(1_000);
+
+        let mut workers = vec![];
+
+        for id in 0..workers_count {
+            workers.push(JinjaEngineWorker::new(
+                id,
+                jinja_engine.clone(),
+                js_channel.clone(),
+                receiver.clone(),
+            ));
+        }
+
+        Self {
+            _workers: workers,
+            workers_rx,
+        }
+    }
+
+    pub fn render(&self, job: JinjaEngineWorkerJob) -> Result<(), CubeError> {
+        self.workers_rx
+            .send_blocking(job)
+            .map_err(|err| CubeError::internal(format!("Unable to schedule rendering: {}", err)))
+    }
+}