From eeea1f27f4e91648f1f4db5213c4812a59bd6180 Mon Sep 17 00:00:00 2001 From: David Michael Barr Date: Sun, 27 Oct 2024 16:49:53 +0900 Subject: [PATCH] Add Lua embedding --- Cargo.lock | 116 +++++++++++++++++++++++++++++ Cargo.toml | 1 + nativelink-error/Cargo.toml | 1 + nativelink-error/src/lib.rs | 17 +++++ nativelink-store/Cargo.toml | 1 + nativelink-store/src/ref_store.rs | 14 ++++ nativelink-util/Cargo.toml | 1 + nativelink-util/src/store_trait.rs | 4 +- src/bin/nativelink.rs | 33 ++++++++ 9 files changed, 187 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index c9609016a..858c993d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -647,6 +647,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bstr" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40723b8fb387abc38f4f4a37c09073622e41dd12327033091ef8950659e6dc0c" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -1317,6 +1327,15 @@ dependencies = [ "digest", ] +[[package]] +name = "home" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +dependencies = [ + "windows-sys 0.52.0", +] + [[package]] name = "http" version = "0.2.12" @@ -1610,6 +1629,25 @@ dependencies = [ "hashbrown 0.15.0", ] +[[package]] +name = "lua-src" +version = "547.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1edaf29e3517b49b8b746701e5648ccb5785cde1c119062cbabbc5d5cd115e42" +dependencies = [ + "cc", +] + +[[package]] +name = "luajit-src" +version = "210.5.11+97813fb" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3015551c284515db7c30c559fc1080f9cb9ee990d1f6fca315451a107c7540bb" +dependencies = [ + "cc", + "which", +] + [[package]] name = "lz4_flex" version = "0.11.3" @@ -1714,6 +1752,50 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "mlua" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f6ddbd668297c46be4bdea6c599dcc1f001a129586272d53170b7ac0a62961e" +dependencies = [ + "bstr", + "either", + "futures-util", + "mlua-sys", + "mlua_derive", + "num-traits", + "parking_lot", + "rustc-hash", +] + +[[package]] +name = "mlua-sys" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9eebac25c35a13285456c88ee2fde93d9aee8bcfdaf03f9d6d12be3391351ec" +dependencies = [ + "cc", + "cfg-if", + "lua-src", + "luajit-src", + "pkg-config", +] + +[[package]] +name = "mlua_derive" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cfc5faa2e0d044b3f5f0879be2920e0a711c97744c42cf1c295cb183668933e" +dependencies = [ + "itertools", + "once_cell", + "proc-macro-error", + "proc-macro2", + "quote", + "regex", + "syn 2.0.79", +] + [[package]] name = "mock_instant" version = "0.5.1" @@ -1737,6 +1819,7 @@ dependencies = [ "hyper 1.4.1", "hyper-util", "mimalloc", + "mlua", "nativelink-config", "nativelink-error", "nativelink-metric", @@ -1781,6 +1864,7 @@ version = "0.5.3" dependencies = [ "fred", "hex", + "mlua", "nativelink-metric", "nativelink-proto", "prost", @@ -1942,6 +2026,7 @@ dependencies = [ "hyper-rustls", "lz4_flex", "memory-stats", + "mlua", "mock_instant", "nativelink-config", "nativelink-error", @@ -1984,6 +2069,7 @@ dependencies = [ "hyper 1.4.1", "hyper-util", "lru", + "mlua", "mock_instant", "nativelink-config", "nativelink-error", @@ -2295,6 +2381,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkg-config" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" + [[package]] name = "powerfmt" version = "0.2.0" @@ -2595,6 +2687,12 @@ version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" +[[package]] +name = "rustc-hash" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "583034fd73374156e66797ed8e5b0d5690409c9226b22d87cb7f19821c05d152" + [[package]] name = "rustc_version" version = "0.4.1" @@ -3522,6 +3620,18 @@ version = "0.25.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1" +[[package]] +name = "which" +version = "6.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ee928febd44d98f2f459a4a79bd4d928591333a494a10a868418ac1b39cf1f" +dependencies = [ + "either", + "home", + "rustix", + "winsafe", +] + [[package]] name = "winapi" version = "0.3.9" @@ -3626,6 +3736,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + [[package]] name = "xmlparser" version = "0.13.6" diff --git a/Cargo.toml b/Cargo.toml index 33fe54582..0d4c8cc77 100644 --- a/Cargo.toml +++ b/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 } diff --git a/nativelink-error/Cargo.toml b/nativelink-error/Cargo.toml index bcbe76e9c..2ce5c646f 100644 --- a/nativelink-error/Cargo.toml +++ b/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 } diff --git a/nativelink-error/src/lib.rs b/nativelink-error/src/lib.rs index b87782c32..955c5e227 100644 --- a/nativelink-error/src/lib.rs +++ b/nativelink-error/src/lib.rs @@ -260,6 +260,23 @@ impl From for tonic::Status { } } +impl From 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 for mlua::Error { + fn from(val: Error) -> Self { + Self::external(val) + } +} + pub trait ResultExt { fn err_tip_with_code(self, tip_fn: F) -> Result where diff --git a/nativelink-store/Cargo.toml b/nativelink-store/Cargo.toml index fa3094417..5ea2dfc85 100644 --- a/nativelink-store/Cargo.toml +++ b/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 } diff --git a/nativelink-store/src/ref_store.rs b/nativelink-store/src/ref_store.rs index d2446d174..eeab1d2f4 100644 --- a/nativelink-store/src/ref_store.rs +++ b/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>(fields: &mut F) { + fields.add_field_method_get("name", |_, this| Ok(this.ref_store_name.clone())); + } + fn add_methods>(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]) + }); + } +} diff --git a/nativelink-util/Cargo.toml b/nativelink-util/Cargo.toml index fb635af60..5fd4bcc4e 100644 --- a/nativelink-util/Cargo.toml +++ b/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 } diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index bf6ca11f6..39b5dbe2b 100644 --- a/nativelink-util/src/store_trait.rs +++ b/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 { diff --git a/src/bin/nativelink.rs b/src/bin/nativelink.rs index 7e2276a63..cd29a91d3 100644 --- a/src/bin/nativelink.rs +++ b/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, } /// 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> { Ok(serde_json5::from_str(&json_contents)?) } +async fn get_lua() -> Result>, Box> { + 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> { init_tracing()?;