Add Lua embedding

Cargo.toml

@@ -53,6 +53,7 @@ futures = { version = "0.3.30", default-features = false }
 hyper = "1.4.1"
 hyper-util = "0.1.9"
 mimalloc = "0.1.43"
+mlua = { version = "0.10.0", features = ["lua54", "vendored", "async", "macros", "send"] }
 parking_lot = "0.12.3"
 rustls-pemfile = { version = "2.2.0", default-features = false }
 scopeguard = { version = "1.2.0", default-features = false }

nativelink-error/Cargo.toml

@@ -14,6 +14,7 @@ fred = { version = "9.2.1", default-features = false, features = [
   "enable-rustls-ring",
 ] }
 hex = { version = "0.4.3", default-features = false }
+mlua = { version = "0.10.0", features = ["lua54", "vendored", "async", "macros"] }
 prost = { version = "0.13.3", default-features = false }
 prost-types = { version = "0.13.3", default-features = false }
 serde = { version = "1.0.210", default-features = false }

nativelink-error/src/lib.rs

@@ -260,6 +260,23 @@ impl From<Error> for tonic::Status {
     }
 }
 
+impl From<mlua::Error> for Error {
+    fn from(val: mlua::Error) -> Self {
+        match val {
+            mlua::Error::CallbackError { traceback, cause } => {
+                Self::new(Code::Internal, traceback).merge(std::sync::Arc::unwrap_or_clone(cause))
+            }
+            _ => Self::new(Code::Internal, val.to_string()),
+        }
+    }
+}
+
+impl From<Error> for mlua::Error {
+    fn from(val: Error) -> Self {
+        Self::external(val)
+    }
+}
+
 pub trait ResultExt<T> {
     fn err_tip_with_code<F, S>(self, tip_fn: F) -> Result<T, Error>
     where

nativelink-store/Cargo.toml

@@ -47,6 +47,7 @@ hyper-rustls = { version = "0.24.2", default-features = false, features = [
   "webpki-roots",
 ] }
 lz4_flex = { version = "0.11.3", default-features = false }
+mlua = { version = "0.10.0", features = ["lua54", "vendored", "async", "macros"] }
 parking_lot = "0.12.3"
 prost = { version = "0.13.3", default-features = false }
 rand = { version = "0.8.5", default-features = false }

nativelink-store/src/ref_store.rs

@@ -151,3 +151,17 @@ impl StoreDriver for RefStore {
 }
 
 default_health_status_indicator!(RefStore);
+
+impl mlua::UserData for RefStore {
+    fn add_fields<F: mlua::UserDataFields<Self>>(fields: &mut F) {
+        fields.add_field_method_get("name", |_, this| Ok(this.ref_store_name.clone()));
+    }
+    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
+        methods.add_async_method("has", |_lua, this, key| async move {
+            let keys = [key];
+            let mut results = [None];
+            this.has_with_results(&keys, &mut results).await?;
+            Ok(results[0])
+        });
+    }
+}

nativelink-util/Cargo.toml

@@ -27,6 +27,7 @@ hex = { version = "0.4.3", default-features = false, features = ["std"] }
 hyper = "1.4.1"
 hyper-util = "0.1.9"
 lru = { version = "0.12.4", default-features = false }
+mlua = { version = "0.10.0", features = ["lua54", "vendored", "async", "macros"] }
 parking_lot = "0.12.3"
 pin-project-lite = "0.2.14"
 prost = { version = "0.13.3", default-features = false }

nativelink-util/src/store_trait.rs

@@ -149,7 +149,7 @@ pub enum StoreOptimizations {
 /// but it can also be a string if the caller wishes to
 /// store the data directly and reference it by a string
 /// directly.
-#[derive(Debug, Eq)]
+#[derive(Debug, Eq, mlua::FromLua)]
 pub enum StoreKey<'a> {
     /// A string key.
     Str(Cow<'a, str>),
@@ -158,6 +158,8 @@ pub enum StoreKey<'a> {
     Digest(DigestInfo),
 }
 
+impl mlua::UserData for StoreKey<'_> {}
+
 impl<'a> StoreKey<'a> {
     /// Creates a new store key from a string.
     pub const fn new_str(s: &'a str) -> Self {

src/bin/nativelink.rs

@@ -110,6 +110,9 @@ struct Args {
     /// Config file to use.
     #[clap(value_parser)]
     config_file: String,
+    /// Lua script to run.
+    #[clap(value_parser)]
+    lua_file: Option<String>,
 }
 
 /// The root metrics collector struct. All metrics will be
@@ -953,6 +956,27 @@ async fn inner_main(
         root_metrics.write().workers = worker_metrics;
     }
 
+    if let Some(script) = futures::executor::block_on(get_lua())? {
+        let lua = mlua::Lua::new();
+        let store_manager = store_manager.clone();
+        let get_store = lua.create_function(move |_, name: String| {
+            Ok(Arc::into_inner(nativelink_store::ref_store::RefStore::new(
+                &nativelink_config::stores::RefStore { name: name.clone() },
+                Arc::downgrade(&store_manager),
+            ))
+            .unwrap())
+        })?;
+        let storekey = lua.create_function(|_, (hash, length): (String, usize)| {
+            let digest = nativelink_util::common::DigestInfo::try_new(&hash, length)?;
+            Ok(nativelink_util::store_trait::StoreKey::Digest(digest))
+        })?;
+        let globals = lua.globals();
+        globals.set("get_store", get_store)?;
+        globals.set("storekey", storekey)?;
+        let f = lua.load(script).into_function()?;
+        let _ = futures::executor::block_on(f.call_async::<()>(()).map_err(Error::from));
+    }
+
     if let Err(e) = try_join_all(root_futures).await {
         panic!("{e:?}");
     };
@@ -969,6 +993,15 @@ async fn get_config() -> Result<CasConfig, Box<dyn std::error::Error>> {
     Ok(serde_json5::from_str(&json_contents)?)
 }
 
+async fn get_lua() -> Result<Option<Vec<u8>>, Box<dyn std::error::Error>> {
+    let args = Args::parse();
+    Ok(if let Some(script) = args.lua_file {
+        Some(std::fs::read(&script).err_tip(|| format!("Could not open Lua script {}", script))?)
+    } else {
+        None
+    })
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     init_tracing()?;