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).
75/// Do not implement this trait manually; define data plugins with
76/// [`define_data_plugin!`](crate::define_data_plugin).
77pub trait DataPlugin: Any {
78    type DataContainer;
79
80    fn init<C: PluginContext>(context: &C) -> Self::DataContainer;
81
82    /// Returns the index into `Context::data_plugins`, the vector of data plugins, where
83    /// the instance of this data plugin can be found.
84    #[must_use]
85    fn index_within_context() -> usize;
86}
87
88#[cfg(test)]
89mod tests {
90    use std::sync::{Arc, Barrier};
91    use std::thread;
92
93    use super::*;
94    use crate::{define_data_plugin, Context};
95
96    // We attempt an out-of-bounds index with a plugin
97    #[test]
98    #[should_panic(
99        expected = "No data plugin found with index = 1000. You must use the `define_data_plugin!` macro to create a data plugin."
100    )]
101    fn test_wrong_data_plugin_impl_index_oob() {
102        // Suppose a user doesn't use the `define_data_plugin` macro and tries to implement it
103        // themselves. What error modes are possible? First lets try an obviously out-of-bounds
104        // index.
105        struct MyDataPlugin;
106        impl DataPlugin for MyDataPlugin {
107            type DataContainer = Vec<u32>;
108
109            fn init<C: PluginContext>(_context: &C) -> Self::DataContainer {
110                vec![]
111            }
112
113            fn index_within_context() -> usize {
114                1000 // arbitrarily out of bounds
115            }
116        }
117
118        let context = Context::new();
119        let _container = context.get_data(MyDataPlugin);
120    }
121
122    // We attempt a collision with a plugin
123    define_data_plugin!(LegitDataPlugin, Vec<u32>, vec![]);
124    #[should_panic(
125        expected = "TypeID does not match data plugin type. You must use the `define_data_plugin!` macro to create a data plugin."
126    )]
127    #[test]
128    fn test_wrong_data_plugin_impl_wrong_type() {
129        // Suppose a user doesn't use the `define_data_plugin` macro and tries
130        // to implement it themselves. What error modes are possible? Here we
131        // test for an index collision.
132        struct MyOtherDataPlugin;
133        impl DataPlugin for MyOtherDataPlugin {
134            type DataContainer = Vec<u8>;
135
136            fn init<C: PluginContext>(_context: &C) -> Self::DataContainer {
137                vec![]
138            }
139
140            fn index_within_context() -> usize {
141                // Several plugins are registered in a test context, so an index of 1 should
142                // collide with another plugin of a different type.
143                LegitDataPlugin::index_within_context()
144            }
145        }
146
147        let context = Context::new();
148        // Make sure the legit plugin is initialized first
149        let _ = context.get_data(LegitDataPlugin);
150
151        // Panics here:
152        let _container = context.get_data(MyOtherDataPlugin);
153        // Some arbitrary code involving `container`
154    }
155
156    // Test thread safety of `initialize_data_plugin_index`.
157    #[test]
158    fn test_multithreaded_plugin_init() {
159        struct DataPluginContainerA;
160        define_data_plugin!(DataPluginA, DataPluginContainerA, DataPluginContainerA);
161        struct DataPluginContainerB;
162        define_data_plugin!(DataPluginB, DataPluginContainerB, DataPluginContainerB);
163        struct DataPluginContainerC;
164        define_data_plugin!(DataPluginC, DataPluginContainerC, DataPluginContainerC);
165        struct DataPluginContainerD;
166        define_data_plugin!(DataPluginD, DataPluginContainerD, DataPluginContainerD);
167
168        // Plugin accessors
169        let accessors: Vec<&(dyn Fn(&Context) + Send + Sync)> = vec![
170            &|ctx: &Context| {
171                let _ = ctx.get_data(DataPluginA);
172            },
173            &|ctx: &Context| {
174                let _ = ctx.get_data(DataPluginB);
175            },
176            &|ctx: &Context| {
177                let _ = ctx.get_data(DataPluginC);
178            },
179            &|ctx: &Context| {
180                let _ = ctx.get_data(DataPluginD);
181            },
182        ];
183
184        let num_threads = 20;
185        let barrier = Arc::new(Barrier::new(num_threads));
186        let mut handles = Vec::with_capacity(num_threads);
187
188        for i in 0..num_threads {
189            let barrier = Arc::clone(&barrier);
190            let accessor = accessors[i % accessors.len()];
191
192            let handle = thread::spawn(move || {
193                let context = Context::new();
194                barrier.wait();
195                accessor(&context);
196            });
197
198            handles.push(handle);
199        }
200
201        for handle in handles {
202            handle.join().expect("Thread panicked");
203        }
204    }
205}
ixa/data_plugin.rs

ixa/
data_plugin.rs
