runtimes/js: add per-env type constructors

Author: eandre

runtimes/js/src/napi_util.rs

@@ -1,4 +1,6 @@
 use napi::{Either, Env, JsFunction, JsObject, JsUnknown};
+use std::collections::HashMap;
+use std::sync::RwLock;
 
 pub trait PromiseHandler: Clone + Send + Sync + 'static {
     type Output: Send + 'static;
@@ -69,3 +71,72 @@ pub fn await_promise<T, H>(
         _ = outer_tx.send(res);
     });
 }
+
+/// EnvMap is a thread-safe map that stores values associated with Env objects.
+/// It is intended for storing one value per napi_env. We need the map to work with
+/// worker pooling, where we can have multiple napi envs that each need their own copy.
+/// It uses a vector under the hood since the number of values is small (one per worker).
+pub struct EnvMap<T> {
+    map: RwLock<Vec<(usize, T)>>,
+}
+
+impl<T> EnvMap<T> {
+    pub const fn new() -> Self {
+        Self {
+            map: RwLock::new(Vec::new()),
+        }
+    }
+
+    pub fn get(&self, env: Env) -> Option<T>
+    where
+        T: Clone,
+    {
+        let elems = self.map.read().unwrap();
+        for (addr, value) in elems.iter() {
+            if *addr == env.raw().addr() {
+                return Some(value.clone());
+            }
+        }
+        None
+    }
+
+    pub fn get_or_init<F>(&self, env: Env, init: F) -> T
+    where
+        T: Clone,
+        F: FnOnce() -> T,
+    {
+        let addr = env.raw().addr();
+
+        // First try to read
+        let num_scanned = {
+            let map = self.map.read().unwrap();
+            for (key, value) in map.iter() {
+                if *key == addr {
+                    return value.clone();
+                }
+            }
+            map.len()
+        };
+
+        // If not found, get write lock and initialize
+        let mut map = self.map.write().unwrap();
+
+        // Double-check in case another thread initialized it.
+        // We only need to check from the last scanned index
+        for (key, value) in map[num_scanned..].iter() {
+            if *key == addr {
+                return value.clone();
+            }
+        }
+
+        let value = init();
+        map.push((addr, value.clone()));
+        value
+    }
+}
+
+impl<T> Default for EnvMap<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}

runtimes/js/src/runtime.rs

@@ -1,6 +1,7 @@
 use crate::api::{new_api_handler, APIRoute, Request};
 use crate::gateway::{Gateway, GatewayConfig};
 use crate::log::Logger;
+use crate::napi_util::EnvMap;
 use crate::pubsub::{PubSubSubscription, PubSubSubscriptionConfig, PubSubTopic};
 use crate::pvalue::{parse_pvalues, transform_pvalues_request, PVals};
 use crate::secret::Secret;
@@ -22,7 +23,7 @@ use std::thread;
 static RUNTIME: OnceLock<napi::Result<Arc<encore_runtime_core::Runtime>>> = OnceLock::new();
 
 // Type constructors registered from javascript so we can create those type from rust
-static TYPE_CONSTRUCTORS: OnceLock<TypeConstructorRefs> = OnceLock::new();
+static TYPE_CONSTRUCTORS: EnvMap<Arc<TypeConstructorRefs>> = EnvMap::new();
 
 struct TypeConstructorRefs {
     decimal: Ref<()>,
@@ -36,7 +37,7 @@ pub struct Runtime {
 #[napi]
 impl Runtime {
     pub fn create_decimal(env: Env, val: &str) -> napi::Result<JsUnknown> {
-        let constructors = TYPE_CONSTRUCTORS.get().ok_or_else(|| {
+        let constructors = TYPE_CONSTRUCTORS.get(env).ok_or_else(|| {
             Error::new(Status::GenericFailure, "Type constructors not initialized")
         })?;
 
@@ -103,10 +104,12 @@ impl Runtime {
             .get_or_init(|| Ok(Arc::new(init_runtime(false)?)))
             .clone()?;
 
-        TYPE_CONSTRUCTORS.get_or_init(|| TypeConstructorRefs {
-            decimal: env
-                .create_reference(options.type_constructors.decimal)
-                .expect("couldn't create reference to Decimal"),
+        TYPE_CONSTRUCTORS.get_or_init(env, || {
+            Arc::new(TypeConstructorRefs {
+                decimal: env
+                    .create_reference(options.type_constructors.decimal)
+                    .expect("couldn't create reference to Decimal"),
+            })
         });
 
         Ok(Self { runtime })