13414194380

Committed 19 Feb 2025 01:52PM UTC coverage: 38.241% (-0.7%) from 38.963%

Build # 13414194380

Build Type

push

github

Committed by

dariusbakunas

Commit Message

feat: add ability to load plugins based on type

Coverage Stats

0 of 38 new or added lines in 4 files covered. (0.0%)

639 of 1671 relevant lines covered (38.24%)

1.17 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/cogrs-core/src/executor/task_queue_manager.rs

use crate::executor::play_iterator::PlayIterator;
use crate::inventory::host::Host;
use crate::inventory::manager::InventoryManager;
use crate::playbook::play::Play;
use crate::strategy::linear::LinearStrategy;
use crate::strategy::Strategy;
use crate::vars::manager::VariableManager;
use anyhow::Result;
use cogrs_plugins::callback::{CallbackPlugin, EventType};
use cogrs_plugins::plugin_type::PluginType;
use serde_json::Value;
use std::cmp::min;
use std::collections::HashMap;
use std::sync::Arc;

pub struct TaskQueueManager {
    forks: usize,
    callbacks_loaded: bool,
    callbacks: HashMap<EventType, Vec<Arc<dyn CallbackPlugin>>>,
    terminated: bool,
    unreachable_hosts: HashMap<String, Host>,
    workers: Vec<tokio::task::JoinHandle<()>>,
}

const DEFAULT_FORKS: usize = 5;

impl TaskQueueManager {
    pub fn new(forks: Option<usize>) -> Self {
        Self {
            callbacks: HashMap::new(),
            callbacks_loaded: false,
            forks: forks.unwrap_or(DEFAULT_FORKS),
            terminated: false,
            unreachable_hosts: HashMap::new(),
            workers: Vec::with_capacity(forks.unwrap_or(DEFAULT_FORKS)),
        }
    }

    pub fn get_worker(&mut self, index: usize) -> Option<&tokio::task::JoinHandle<()>> {
        self.workers.get(index)
    }

    pub fn set_worker(&mut self, index: usize, worker: tokio::task::JoinHandle<()>) {
        self.workers.insert(index, worker);
    }

    /// Iterates over the roles/tasks in a play, using the given (or default)
    /// strategy for queueing tasks. The default is the linear strategy, which
    /// operates like classic Ansible by keeping all hosts in lock-step with
    /// a given task (meaning no hosts move on to the next task until all hosts
    /// are done with the current task).
    pub async fn run(
        &mut self,
        play: Play,
        variable_manager: &VariableManager,
        inventory_manager: &InventoryManager,
    ) -> Result<()> {
        self.load_callbacks(
            // TODO: add logic to get callback plugin path
            "/Users/darius/Programming/cogrs/dist/minimal-apple_x86_64-apple-darwin",
        );
        let all_vars = variable_manager.get_vars(Some(&play), None, None, true, true);

        self.emit_event(EventType::PlaybookOnPlayStart, None).await;

        let strategy = *play.strategy();

        let mut play_iterator = PlayIterator::new(play);
        play_iterator.init(inventory_manager)?;

        self.forks = min(self.forks, play_iterator.batch_size());

        match strategy {
            Strategy::Linear => {
                let mut strategy = LinearStrategy::new(self, inventory_manager, variable_manager);
                strategy.run(&mut play_iterator).await?;
            }
            Strategy::Free => {
                todo!()
            }
        }

        Ok(())
    }

    pub fn register_callback(&mut self, callback: Box<dyn CallbackPlugin>) {
        let callback: Arc<dyn CallbackPlugin> = Arc::from(callback);
        for event in callback.get_interested_events() {
            self.callbacks
                .entry(event)
                .or_insert_with(Vec::new)
                .push(callback.clone());
        }
    }

    pub fn get_unreachable_hosts(&self) -> &HashMap<String, Host> {
        &self.unreachable_hosts
    }

    pub fn is_terminated(&self) -> bool {
        self.terminated
    }

    fn load_callbacks(&mut self, plugin_dir: &str) -> Result<()> {
        // TODO: move this to plugin loader
        use libloading::{Library, Symbol};
        use std::fs;

        let plugin_extension = if cfg!(target_os = "windows") {
            "dll"
        } else if cfg!(target_os = "macos") {
            "dylib"
        } else {
            "so"
        };

        if self.callbacks_loaded {
            return Ok(());
        }

        for entry in fs::read_dir(plugin_dir).expect("Invalid plugin directory") {
            let path = entry.expect("Failed to read entry").path();
            if path.extension().and_then(|e| e.to_str()) == Some(plugin_extension) {
                unsafe {
                    let lib = Library::new(&path).expect("Failed to load plugin");

                    let plugin_type_fn: Symbol<unsafe extern "C" fn() -> u64> =
                        lib.get(b"plugin_type").unwrap();
                    let plugin_type_value = plugin_type_fn();

                    // Decode the numeric value into the PluginType enum
                    let plugin_type = match PluginType::from_u64(plugin_type_value) {
                        Some(pt) => pt,
                        None => {
                            println!("Skipping unknown plugin type {}", plugin_type_value);
                            continue; // Skip loading this plugin
                        }
                    };

                    match plugin_type {
                        PluginType::Callback => {
                            let create_callback: Symbol<fn() -> Box<dyn CallbackPlugin>> =
                                match lib.get(b"create_plugin") {
                                    Ok(symbol) => symbol,
                                    Err(_) => {
                                        println!(
                                        "Skipping plugin {:?}: missing `create_plugin` function.",
                                        path
                                    );
                                        continue;
                                    }
                                };

                            let plugin = create_callback();

                            // Register the plugin for events
                            self.register_callback(plugin);
                        }
                        _ => {
                            println!("Skipping OtherPlugin: Not supported.");
                        }
                    }
                }
            }
        }

        self.callbacks_loaded = true;
        Ok(())
    }

    pub async fn emit_event(&self, event: EventType, data: Option<Value>) {
        if let Some(callbacks) = self.callbacks.get(&event) {
            // Spawn and collect tasks
            let tasks: Vec<_> = callbacks
                .iter()
                .map(|callback| {
                    let callback = callback.clone();
                    let event = event.clone();
                    let data = data.clone();
                    tokio::spawn(async move {
                        callback.on_event(&event, data.as_ref()); // Async invocation
                    })
                })
                .collect();

            // Wait for all spawned tasks to complete
            for task in tasks {
                if let Err(err) = task.await {
                    eprintln!("Callback task panicked: {:?}", err);
                }
            }
        }
    }

    pub fn forks(&self) -> usize {
        self.forks
    }
}

1	use crate::executor::play_iterator::PlayIterator;
2	use crate::inventory::host::Host;
3	use crate::inventory::manager::InventoryManager;
4	use crate::playbook::play::Play;
5	use crate::strategy::linear::LinearStrategy;
6	use crate::strategy::Strategy;
7	use crate::vars::manager::VariableManager;
8	use anyhow::Result;
9	use cogrs_plugins::callback::{CallbackPlugin, EventType};
10	use cogrs_plugins::plugin_type::PluginType;
11	use serde_json::Value;
12	use std::cmp::min;
13	use std::collections::HashMap;
14	use std::sync::Arc;
15
16	pub struct TaskQueueManager {
17	forks: usize,
18	callbacks_loaded: bool,
19	callbacks: HashMap<EventType, Vec<Arc<dyn CallbackPlugin>>>,
20	terminated: bool,
21	unreachable_hosts: HashMap<String, Host>,
22	workers: Vec<tokio::task::JoinHandle<()>>,
23	}
24
25	const DEFAULT_FORKS: usize = 5;
26
27	impl TaskQueueManager {
28	pub fn new(forks: Option<usize>) -> Self {	×
29	Self {
30	callbacks: HashMap::new(),	×
31	callbacks_loaded: false,
32	forks: forks.unwrap_or(DEFAULT_FORKS),	×
33	terminated: false,
34	unreachable_hosts: HashMap::new(),	×
35	workers: Vec::with_capacity(forks.unwrap_or(DEFAULT_FORKS)),	×
36	}
37	}
38
39	pub fn get_worker(&mut self, index: usize) -> Option<&tokio::task::JoinHandle<()>> {	×
40	self.workers.get(index)	×
41	}
42
43	pub fn set_worker(&mut self, index: usize, worker: tokio::task::JoinHandle<()>) {	×
44	self.workers.insert(index, worker);	×
45	}
46
47	/// Iterates over the roles/tasks in a play, using the given (or default)
48	/// strategy for queueing tasks. The default is the linear strategy, which
49	/// operates like classic Ansible by keeping all hosts in lock-step with
50	/// a given task (meaning no hosts move on to the next task until all hosts
51	/// are done with the current task).
52	pub async fn run(	×
53	&mut self,
54	play: Play,
55	variable_manager: &VariableManager,
56	inventory_manager: &InventoryManager,
57	) -> Result<()> {
58	self.load_callbacks(	×
59	// TODO: add logic to get callback plugin path
60	"/Users/darius/Programming/cogrs/dist/minimal-apple_x86_64-apple-darwin",
61	);
62	let all_vars = variable_manager.get_vars(Some(&play), None, None, true, true);	×
63
64	self.emit_event(EventType::PlaybookOnPlayStart, None).await;	×
65
66	let strategy = *play.strategy();	×
67
68	let mut play_iterator = PlayIterator::new(play);	×
69	play_iterator.init(inventory_manager)?;	×
70
NEW 71	self.forks = min(self.forks, play_iterator.batch_size());	×
72
73	match strategy {	×
74	Strategy::Linear => {
75	let mut strategy = LinearStrategy::new(self, inventory_manager, variable_manager);	×
76	strategy.run(&mut play_iterator).await?;	×
77	}
78	Strategy::Free => {
79	todo!()
80	}
81	}
82
83	Ok(())	×
84	}
85
86	pub fn register_callback(&mut self, callback: Box<dyn CallbackPlugin>) {	×
87	let callback: Arc<dyn CallbackPlugin> = Arc::from(callback);	×
88	for event in callback.get_interested_events() {	×
89	self.callbacks	×
90	.entry(event)
91	.or_insert_with(Vec::new)
92	.push(callback.clone());	×
93	}
94	}
95
96	pub fn get_unreachable_hosts(&self) -> &HashMap<String, Host> {	×
97	&self.unreachable_hosts	×
98	}
99
100	pub fn is_terminated(&self) -> bool {	×
101	self.terminated	×
102	}
103
NEW 104	fn load_callbacks(&mut self, plugin_dir: &str) -> Result<()> {	×
105	// TODO: move this to plugin loader
106	use libloading::{Library, Symbol};
107	use std::fs;
108
109	let plugin_extension = if cfg!(target_os = "windows") {
110	"dll"
111	} else if cfg!(target_os = "macos") {
112	"dylib"
113	} else {
114	"so"	×
115	};
116
117	if self.callbacks_loaded {	×
NEW 118	return Ok(());	×
119	}
120
121	for entry in fs::read_dir(plugin_dir).expect("Invalid plugin directory") {	×
122	let path = entry.expect("Failed to read entry").path();	×
123	if path.extension().and_then(\|e\| e.to_str()) == Some(plugin_extension) {	×
124	unsafe {
125	let lib = Library::new(&path).expect("Failed to load plugin");	×
126
NEW 127	let plugin_type_fn: Symbol<unsafe extern "C" fn() -> u64> =	×
128	lib.get(b"plugin_type").unwrap();
NEW 129	let plugin_type_value = plugin_type_fn();	×
130
131	// Decode the numeric value into the PluginType enum
NEW 132	let plugin_type = match PluginType::from_u64(plugin_type_value) {	×
NEW 133	Some(pt) => pt,	×
134	None => {
NEW 135	println!("Skipping unknown plugin type {}", plugin_type_value);	×
136	continue; // Skip loading this plugin
137	}
138	};
139
NEW 140	match plugin_type {	×
141	PluginType::Callback => {
NEW 142	let create_callback: Symbol<fn() -> Box<dyn CallbackPlugin>> =	×
143	match lib.get(b"create_plugin") {
NEW 144	Ok(symbol) => symbol,	×
145	Err(_) => {
NEW 146	println!(	×
147	"Skipping plugin {:?}: missing `create_plugin` function.",
148	path
149	);
150	continue;
151	}
152	};
153
NEW 154	let plugin = create_callback();	×
155
156	// Register the plugin for events
NEW 157	self.register_callback(plugin);	×
158	}
159	_ => {
NEW 160	println!("Skipping OtherPlugin: Not supported.");	×
161	}
162	}
163	}
164	}
165	}
166
NEW 167	self.callbacks_loaded = true;	×
NEW 168	Ok(())	×
169	}
170
171	pub async fn emit_event(&self, event: EventType, data: Option<Value>) {	×
172	if let Some(callbacks) = self.callbacks.get(&event) {	×
173	// Spawn and collect tasks
174	let tasks: Vec<_> = callbacks	×
175	.iter()
176	.map(\|callback\| {	×
177	let callback = callback.clone();	×
178	let event = event.clone();	×
179	let data = data.clone();	×
180	tokio::spawn(async move {	×
181	callback.on_event(&event, data.as_ref()); // Async invocation	×
182	})
183	})
184	.collect();
185
186	// Wait for all spawned tasks to complete
187	for task in tasks {	×
188	if let Err(err) = task.await {	×
189	eprintln!("Callback task panicked: {:?}", err);	×
190	}
191	}
192	}
193	}
194
195	pub fn forks(&self) -> usize {	×
196	self.forks	×
197	}
198	}

dariusbakunas / cogrs / 13414194380

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous