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

foundationdb_simulation/
lib.rs
