foundationdb_simulation/
lib.rs

1#![warn(missing_docs)]
2#![doc = include_str!("../README.md")]
3
4use std::{ptr::NonNull, sync::Arc};
5
6use foundationdb::Database;
7use foundationdb_sys::FDBDatabase as FDBDatabaseAlias;
8
9mod bindings;
10mod fdb_rt;
11
12use bindings::{FDBDatabase, FDBMetrics, FDBPromise, FDBWorkload, OpaqueWorkload, Promise};
13pub use bindings::{Metric, Metrics, Severity, WorkloadContext};
14use fdb_rt::fdb_spawn;
15
16// -----------------------------------------------------------------------------
17// User friendly types
18
19/// Rust representation of a simulated FoundationDB database
20pub type SimDatabase = Arc<Database>;
21/// Rust representation of a FoundationDB workload
22pub type WrappedWorkload = FDBWorkload;
23
24/// Equivalent to the C++ abstract class `FDBWorkload`
25#[allow(async_fn_in_trait)]
26pub trait RustWorkload {
27    /// This method is called by the tester during the setup phase.
28    /// It should be used to populate the database.
29    ///
30    /// # Arguments
31    ///
32    /// * `db` - The simulated database.
33    async fn setup(&mut self, db: SimDatabase);
34
35    /// This method should run the actual test.
36    ///
37    /// # Arguments
38    ///
39    /// * `db` - The simulated database.
40    async fn start(&mut self, db: SimDatabase);
41
42    /// This method is called when the tester completes.
43    /// A workload should run any consistency/correctness tests during this phase.
44    ///
45    /// # Arguments
46    ///
47    /// * `db` - The simulated database.
48    async fn check(&mut self, db: SimDatabase);
49
50    /// If a workload collects metrics (like latencies or throughput numbers), these should be reported back here.
51    /// The multitester (or test orchestrator) will collect all metrics from all test clients and it will aggregate them.
52    ///
53    /// # Arguments
54    ///
55    /// * `out` - A metric sink
56    fn get_metrics(&self, out: Metrics);
57
58    /// Set the check timeout in simulated seconds for this workload.
59    fn get_check_timeout(&self) -> f64;
60}
61
62/// Equivalent to the C++ abstract class `FDBWorkloadFactory`
63pub trait RustWorkloadFactory {
64    /// The runtime FDB_API_VERSION to use
65    const FDB_API_VERSION: u32 = foundationdb_sys::FDB_API_VERSION;
66    /// If the test file contains a key-value pair workloadName the value will be passed to this method (empty string otherwise).
67    /// This way, a library author can implement many workloads in one library and use the test file to chose which one to run
68    /// (or run multiple workloads either concurrently or serially).
69    fn create(name: String, context: WorkloadContext) -> WrappedWorkload;
70}
71
72/// Automatically implements a WorkloadFactory for a single workload
73pub trait SingleRustWorkload: RustWorkload {
74    /// The runtime FDB_API_VERSION to use
75    const FDB_API_VERSION: u32 = foundationdb_sys::FDB_API_VERSION;
76    /// The implicit WorkloadFactory will call this method uppon each instantiation
77    fn new(name: String, context: WorkloadContext) -> Self;
78}
79
80// -----------------------------------------------------------------------------
81// C to Rust bindings
82
83fn check_database_ref(database: SimDatabase) {
84    if Arc::strong_count(&database) != 1 || Arc::weak_count(&database) != 0 {
85        eprintln!("Reference to Database kept between phases (setup/start/check). All references should be dropped.");
86        std::process::exit(1);
87    }
88    std::mem::forget(database);
89}
90
91unsafe fn database_new(raw_database: *mut FDBDatabase) -> SimDatabase {
92    Arc::new(Database::new_from_pointer(NonNull::new_unchecked(
93        raw_database as *mut FDBDatabaseAlias,
94    )))
95}
96unsafe extern "C" fn workload_setup<W: RustWorkload + 'static>(
97    raw_workload: *mut OpaqueWorkload,
98    raw_database: *mut FDBDatabase,
99    raw_promise: FDBPromise,
100) {
101    let workload = &mut *(raw_workload as *mut W);
102    let database = database_new(raw_database);
103    let done = Promise::new(raw_promise);
104    fdb_spawn(async move {
105        workload.setup(database.clone()).await;
106        check_database_ref(database);
107        done.send(true);
108    });
109}
110unsafe extern "C" fn workload_start<W: RustWorkload + 'static>(
111    raw_workload: *mut OpaqueWorkload,
112    raw_database: *mut FDBDatabase,
113    raw_promise: FDBPromise,
114) {
115    let workload = &mut *(raw_workload as *mut W);
116    let database = database_new(raw_database);
117    let done = Promise::new(raw_promise);
118    fdb_spawn(async move {
119        workload.start(database.clone()).await;
120        check_database_ref(database);
121        done.send(true);
122    });
123}
124unsafe extern "C" fn workload_check<W: RustWorkload + 'static>(
125    raw_workload: *mut OpaqueWorkload,
126    raw_database: *mut FDBDatabase,
127    raw_promise: FDBPromise,
128) {
129    let workload = &mut *(raw_workload as *mut W);
130    let database = database_new(raw_database);
131    let done = Promise::new(raw_promise);
132    fdb_spawn(async move {
133        workload.check(database.clone()).await;
134        check_database_ref(database);
135        done.send(true);
136    });
137}
138unsafe extern "C" fn workload_get_metrics<W: RustWorkload>(
139    raw_workload: *mut OpaqueWorkload,
140    raw_metrics: FDBMetrics,
141) {
142    let workload = &*(raw_workload as *mut W);
143    let out = Metrics::new(raw_metrics);
144    workload.get_metrics(out)
145}
146unsafe extern "C" fn workload_get_check_timeout<W: RustWorkload>(
147    raw_workload: *mut OpaqueWorkload,
148) -> f64 {
149    let workload = &*(raw_workload as *mut W);
150    workload.get_check_timeout()
151}
152unsafe extern "C" fn workload_drop<W: RustWorkload>(raw_workload: *mut OpaqueWorkload) {
153    unsafe { drop(Box::from_raw(raw_workload as *mut W)) };
154}
155
156impl WrappedWorkload {
157    pub fn new<W: RustWorkload + 'static>(workload: W) -> Self {
158        let workload = Box::into_raw(Box::new(workload));
159        WrappedWorkload {
160            inner: workload as *mut _,
161            setup: Some(workload_setup::<W>),
162            start: Some(workload_start::<W>),
163            check: Some(workload_check::<W>),
164            getMetrics: Some(workload_get_metrics::<W>),
165            getCheckTimeout: Some(workload_get_check_timeout::<W>),
166            free: Some(workload_drop::<W>),
167        }
168    }
169}
170// -----------------------------------------------------------------------------
171// Registration hooks
172
173#[doc(hidden)]
174/// Primitives exposed for the registrations hooks, should not be used otherwise
175pub mod internals {
176    pub use crate::bindings::{str_from_c, FDBWorkloadContext};
177
178    #[cfg(feature = "cpp-abi")]
179    extern "C" {
180        pub fn workloadCppFactory(logger: *const u8) -> *const u8;
181    }
182
183    #[allow(non_snake_case)]
184    #[cfg(not(feature = "cpp-abi"))]
185    pub unsafe extern "C" fn workloadCppFactory(_logger: *const u8) -> *const u8 {
186        eprintln!(
187            "This Rust workload was compiled without the C++ shim adapter. To fix this, either:
188
189- Re-run the simulation with `useCAPI = true` (FoundationDB 7.4 or newer), or
190- Recompile the workload with FoundationDB versions prior to 7.4 or the `cpp-abi` feature"
191        );
192        std::process::exit(1);
193    }
194}
195
196/// Register a [RustWorkloadFactory].
197/// /!\ Should be called only once.
198#[macro_export]
199macro_rules! register_factory {
200    ($name:ident) => {
201        #[no_mangle]
202        extern "C" fn workloadCFactory(
203            raw_name: *const i8,
204            raw_context: $crate::internals::FDBWorkloadContext,
205        ) -> $crate::WrappedWorkload {
206            use std::sync::atomic::{AtomicBool, Ordering};
207            static DONE: AtomicBool = AtomicBool::new(false);
208            if DONE
209                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
210                .is_ok()
211            {
212                let version = <$name as $crate::RustWorkloadFactory>::FDB_API_VERSION;
213                let _ = foundationdb::api::FdbApiBuilder::default()
214                    .set_runtime_version(version as i32)
215                    .build();
216                println!("FDB API version selected: {version}");
217            }
218            let name = $crate::internals::str_from_c(raw_name);
219            let context = $crate::WorkloadContext::new(raw_context);
220            <$name as $crate::RustWorkloadFactory>::create(name, context)
221        }
222        #[no_mangle]
223        extern "C" fn workloadFactory(logger: *const u8) -> *const u8 {
224            unsafe { $crate::internals::workloadCppFactory(logger) }
225        }
226    };
227}
228
229/// Register a [SingleRustWorkload] and creates an implicit WorkloadFactory.
230/// /!\ Should be called only once.
231#[macro_export]
232macro_rules! register_workload {
233    ($name:ident) => {
234        #[no_mangle]
235        extern "C" fn workloadCFactory(
236            raw_name: *const i8,
237            raw_context: $crate::internals::FDBWorkloadContext,
238        ) -> $crate::WrappedWorkload {
239            use std::sync::atomic::{AtomicBool, Ordering};
240            static DONE: AtomicBool = AtomicBool::new(false);
241            if DONE
242                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
243                .is_ok()
244            {
245                let version = <$name as $crate::SingleRustWorkload>::FDB_API_VERSION;
246                let _ = foundationdb::api::FdbApiBuilder::default()
247                    .set_runtime_version(version as i32)
248                    .build();
249                println!("FDB API version selected: {version}");
250            }
251            let name = $crate::internals::str_from_c(raw_name);
252            let context = $crate::WorkloadContext::new(raw_context);
253            $crate::WrappedWorkload::new(<$name as $crate::SingleRustWorkload>::new(name, context))
254        }
255        #[no_mangle]
256        extern "C" fn workloadFactory(logger: *const u8) -> *const u8 {
257            unsafe { $crate::internals::workloadCppFactory(logger) }
258        }
259    };
260}
foundationdb_simulation/lib.rs

foundationdb_simulation/
lib.rs
