From eb83ec4edc733fee34b968a406bce86be620b4e2 Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Thu, 20 Jun 2024 16:26:47 +0000
Subject: [PATCH] Make LeanDefinitionalEngine thread safe (#227)

Signed-off-by: Andrew Wells <anmwells@amazon.com>
---
 cedar-drt/Cargo.toml         |  2 +-
 cedar-drt/src/lean_impl.rs   | 16 +++++++++++-----
 cedar-drt/tests/benchmark.rs |  3 ---
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/cedar-drt/Cargo.toml b/cedar-drt/Cargo.toml
index b6a08f827..b6f20e9ab 100644
--- a/cedar-drt/Cargo.toml
+++ b/cedar-drt/Cargo.toml
@@ -13,7 +13,7 @@ cedar-policy-core = { path = "../cedar/cedar-policy-core", version = "4.*", feat
 cedar-policy-validator = { path = "../cedar/cedar-policy-validator", version = "4.*", features = ["arbitrary"] }
 cedar-policy-formatter = { path = "../cedar/cedar-policy-formatter", version = "4.*" }
 cedar-testing = { path = "../cedar/cedar-testing", version = "4.*" }
-lean-sys = { version = "0.0.5", features = ["small_allocator"], default-features = false }
+lean-sys = { version = "0.0.7", features = ["small_allocator"], default-features = false }
 miette = "7.1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
diff --git a/cedar-drt/src/lean_impl.rs b/cedar-drt/src/lean_impl.rs
index dbea58ee9..d7faf3cdb 100644
--- a/cedar-drt/src/lean_impl.rs
+++ b/cedar-drt/src/lean_impl.rs
@@ -34,11 +34,8 @@ use cedar_testing::cedar_test_impl::*;
 pub use lean_sys::init::lean_initialize;
 pub use lean_sys::lean_object;
 pub use lean_sys::string::lean_mk_string;
-use lean_sys::{lean_dec, lean_dec_ref, lean_io_result_is_ok, lean_io_result_show_error};
-use lean_sys::{
-    lean_initialize_runtime_module_locked, lean_io_mark_end_initialization, lean_io_mk_world,
-    lean_string_cstr,
-};
+use lean_sys::{lean_dec, lean_dec_ref, lean_io_result_is_ok, lean_io_result_show_error, lean_initialize_runtime_module_locked, lean_io_mark_end_initialization, lean_io_mk_world,
+    lean_string_cstr, lean_initialize_thread, lean_finalize_thread};
 use log::info;
 use miette::miette;
 use serde::Deserialize;
@@ -150,6 +147,7 @@ impl LeanDefinitionalEngine {
                 lean_io_mark_end_initialization();
             };
         });
+        unsafe{lean_initialize_thread()};
         Self {}
     }
 
@@ -374,6 +372,14 @@ impl LeanDefinitionalEngine {
     }
 }
 
+impl Drop for LeanDefinitionalEngine {
+    fn drop(&mut self) {
+        unsafe {
+            lean_finalize_thread()
+        }
+    }
+}
+
 impl CedarTestImplementation for LeanDefinitionalEngine {
     fn is_authorized(
         &self,
diff --git a/cedar-drt/tests/benchmark.rs b/cedar-drt/tests/benchmark.rs
index 5cc6444d0..2bb2bd1d9 100644
--- a/cedar-drt/tests/benchmark.rs
+++ b/cedar-drt/tests/benchmark.rs
@@ -163,10 +163,7 @@ fn print_summary(auth_times: HashMap<&str, Vec<f64>>, val_times: HashMap<&str, V
     }
 }
 
-// Currently, running this in conjunction with existing tests will cause an error (#227).
-// In order see the printed output from this test, run `cargo test -- --ignored --nocapture`.
 #[test]
-#[ignore = "Can only run one Lean FFI thread currently."]
 fn run_all_tests() {
     let rust_impl = RustEngine::new();
     let lean_impl = LeanDefinitionalEngine::new();