ixa/
data_plugin.rs

1use std::any::{Any, TypeId};
2use std::cell::RefCell;
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{LazyLock, Mutex};
5
6use crate::{HashSet, PluginContext};
7
8/// A collection of [`TypeId`]s of all [`DataPlugin`] types linked into the code.
9static DATA_PLUGINS: LazyLock<Mutex<RefCell<HashSet<TypeId>>>> =
10    LazyLock::new(|| Mutex::new(RefCell::new(HashSet::default())));
11
12pub fn add_data_plugin_to_registry<T: DataPlugin>() {
13    DATA_PLUGINS
14        .lock()
15        .unwrap()
16        .borrow_mut()
17        .insert(TypeId::of::<T>());
18}
19
20#[must_use]
21pub fn get_data_plugin_ids() -> Vec<TypeId> {
22    DATA_PLUGINS
23        .lock()
24        .unwrap()
25        .borrow()
26        .iter()
27        .copied()
28        .collect()
29}
30
31#[must_use]
32pub fn get_data_plugin_count() -> usize {
33    DATA_PLUGINS.lock().unwrap().borrow().len()
34}
35
36/// Global data plugin index counter, keeps track of the index that will be assigned to the next
37/// data plugin that requests an index.
38///
39/// Instead of storing data plugins in a [`HashMap`] in [`Context`], we store them in a vector. To fetch
40/// the data plugin, we ask the data plugin type for the index into [`Context::data_plugins`] at
41/// which an instance of the data plugin type should be stored. Accessing a data plugin, then, is
42/// just an index into an array.
43static NEXT_DATA_PLUGIN_INDEX: Mutex<usize> = Mutex::new(0);
44
45/// Acquires a global lock on the next available plugin index, but only increments it if we
46/// successfully initialize the provided index. (Must be `pub`, as it's called from within a macro.)
47#[must_use]
48pub fn initialize_data_plugin_index(plugin_index: &AtomicUsize) -> usize {
49    // Acquire a global lock.
50    let mut guard = NEXT_DATA_PLUGIN_INDEX.lock().unwrap();
51    let candidate = *guard;
52
53    // Try to claim the candidate index. Here we guard against the potential race condition that
54    // another instance of this plugin in another thread just initialized the index prior to us
55    // obtaining the lock. If the index has been initialized beneath us, we do not update
56    // `NEXT_DATA_PLUGIN_INDEX`, we just return the value `plugin_index` was initialized to.
57    // For a justification of the data ordering, see:
58    //     https://github.com/CDCgov/ixa/pull/477#discussion_r2244302872
59    match plugin_index.compare_exchange(usize::MAX, candidate, Ordering::AcqRel, Ordering::Acquire)
60    {
61        Ok(_) => {
62            // We won the race — increment the global next plugin index and return the new index
63            *guard += 1;
64            candidate
65        }
66        Err(existing) => {
67            // Another thread beat us — don’t increment the global next plugin index,
68            // just return existing
69            existing
70        }
71    }
72}
73
74/// A trait for objects that can provide data containers to be held by [`Context`](crate::Context)
75pub trait DataPlugin: Any {
76    type DataContainer;
77
78    fn init<C: PluginContext>(context: &C) -> Self::DataContainer;
79
80    /// Returns the index into `Context::data_plugins`, the vector of data plugins, where
81    /// the instance of this data plugin can be found.
82    #[must_use]
83    fn index_within_context() -> usize;
84}
85
86#[cfg(test)]
87mod tests {
88    use std::sync::{Arc, Barrier};
89    use std::thread;
90
91    use super::*;
92    use crate::{define_data_plugin, Context};
93
94    // We attempt an out-of-bounds index with a plugin
95    #[test]
96    #[should_panic(
97        expected = "No data plugin found with index = 1000. You must use the `define_data_plugin!` macro to create a data plugin."
98    )]
99    fn test_wrong_data_plugin_impl_index_oob() {
100        // Suppose a user doesn't use the `define_data_plugin` macro and tries to implement it
101        // themselves. What error modes are possible? First lets try an obviously out-of-bounds
102        // index.
103        struct MyDataPlugin;
104        impl DataPlugin for MyDataPlugin {
105            type DataContainer = Vec<u32>;
106
107            fn init<C: PluginContext>(_context: &C) -> Self::DataContainer {
108                vec![]
109            }
110
111            fn index_within_context() -> usize {
112                1000 // arbitrarily out of bounds
113            }
114        }
115
116        let context = Context::new();
117        let _container = context.get_data(MyDataPlugin);
118    }
119
120    // We attempt a collision with a plugin
121    define_data_plugin!(LegitDataPlugin, Vec<u32>, vec![]);
122    #[should_panic(
123        expected = "TypeID does not match data plugin type. You must use the `define_data_plugin!` macro to create a data plugin."
124    )]
125    #[test]
126    fn test_wrong_data_plugin_impl_wrong_type() {
127        // Suppose a user doesn't use the `define_data_plugin` macro and tries
128        // to implement it themselves. What error modes are possible? Here we
129        // test for an index collision.
130        struct MyOtherDataPlugin;
131        impl DataPlugin for MyOtherDataPlugin {
132            type DataContainer = Vec<u8>;
133
134            fn init<C: PluginContext>(_context: &C) -> Self::DataContainer {
135                vec![]
136            }
137
138            fn index_within_context() -> usize {
139                // Several plugins are registered in a test context, so an index of 1 should
140                // collide with another plugin of a different type.
141                LegitDataPlugin::index_within_context()
142            }
143        }
144
145        let context = Context::new();
146        // Make sure the legit plugin is initialized first
147        let _ = context.get_data(LegitDataPlugin);
148
149        // Panics here:
150        let _container = context.get_data(MyOtherDataPlugin);
151        // Some arbitrary code involving `container`
152    }
153
154    // Test thread safety of `initialize_data_plugin_index`.
155    #[test]
156    fn test_multithreaded_plugin_init() {
157        struct DataPluginContainerA;
158        define_data_plugin!(DataPluginA, DataPluginContainerA, DataPluginContainerA);
159        struct DataPluginContainerB;
160        define_data_plugin!(DataPluginB, DataPluginContainerB, DataPluginContainerB);
161        struct DataPluginContainerC;
162        define_data_plugin!(DataPluginC, DataPluginContainerC, DataPluginContainerC);
163        struct DataPluginContainerD;
164        define_data_plugin!(DataPluginD, DataPluginContainerD, DataPluginContainerD);
165
166        // Plugin accessors
167        let accessors: Vec<&(dyn Fn(&Context) + Send + Sync)> = vec![
168            &|ctx: &Context| {
169                let _ = ctx.get_data(DataPluginA);
170            },
171            &|ctx: &Context| {
172                let _ = ctx.get_data(DataPluginB);
173            },
174            &|ctx: &Context| {
175                let _ = ctx.get_data(DataPluginC);
176            },
177            &|ctx: &Context| {
178                let _ = ctx.get_data(DataPluginD);
179            },
180        ];
181
182        let num_threads = 20;
183        let barrier = Arc::new(Barrier::new(num_threads));
184        let mut handles = Vec::with_capacity(num_threads);
185
186        for i in 0..num_threads {
187            let barrier = Arc::clone(&barrier);
188            let accessor = accessors[i % accessors.len()];
189
190            let handle = thread::spawn(move || {
191                let context = Context::new();
192                barrier.wait();
193                accessor(&context);
194            });
195
196            handles.push(handle);
197        }
198
199        for handle in handles {
200            handle.join().expect("Thread panicked");
201        }
202    }
203}
ixa/data_plugin.rs

ixa/
data_plugin.rs
