13381193139

Committed 18 Feb 2025 01:09AM UTC coverage: 39.987% (-1.2%) from 41.226%

Build # 13381193139

Build Type

push

github

Committed by

dariusbakunas

Commit Message

refactor: cleanup

Coverage Stats

1 of 1 new or added line in 1 file covered. (100.0%)

182 existing lines in 6 files now uncovered.

639 of 1598 relevant lines covered (39.99%)

1.32 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 serde_json::Value;
use std::cmp::min;
use std::collections::HashMap;
use std::sync::Arc;

pub struct TaskQueueManager<'a> {
    forks: u32,
    callbacks_loaded: bool,
    inventory_manager: &'a InventoryManager,
    variable_manager: &'a VariableManager<'a>,
    callbacks: HashMap<EventType, Vec<Arc<dyn CallbackPlugin>>>,
    terminated: bool,
    unreachable_hosts: HashMap<String, Host>,
}

const DEFAULT_FORKS: u32 = 5;

impl<'a> TaskQueueManager<'a> {
    pub fn new(
        forks: Option<u32>,
        inventory_manager: &'a InventoryManager,
        variable_manager: &'a VariableManager,
    ) -> Self {
        Self {
            callbacks: HashMap::new(),
            callbacks_loaded: false,
            forks: forks.unwrap_or(DEFAULT_FORKS),
            inventory_manager,
            variable_manager,
            terminated: false,
            unreachable_hosts: HashMap::new(),
        }
    }

    pub fn inventory_manager(&self) -> &InventoryManager {
        self.inventory_manager
    }

    pub fn variable_manager(&self) -> &VariableManager {
        self.variable_manager
    }

    pub async fn run(&mut self, play: &Play) -> 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 = self
            .variable_manager
            .get_vars(Some(play), None, None, true, true);

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

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

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

        match play.strategy() {
            Strategy::Linear => {
                let mut strategy = LinearStrategy::new(&self);
                strategy.run(&mut play_iterator)?;
            }
            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) {
        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;
        }

        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");

                    // Dynamically load the callback creation function
                    let create_callback: Symbol<fn() -> Box<dyn CallbackPlugin>> = lib
                        .get(b"create_plugin")
                        .expect("Failed to find create_plugin function");

                    let plugin = create_callback();

                    // Register the plugin for events
                    self.register_callback(plugin);
                }
            }
        }

        self.callbacks_loaded = true
    }

    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);
                }
            }
        }
    }
}

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 serde_json::Value;
11	use std::cmp::min;
12	use std::collections::HashMap;
13	use std::sync::Arc;
14
15	pub struct TaskQueueManager<'a> {
16	forks: u32,
17	callbacks_loaded: bool,
18	inventory_manager: &'a InventoryManager,
19	variable_manager: &'a VariableManager<'a>,
20	callbacks: HashMap<EventType, Vec<Arc<dyn CallbackPlugin>>>,
21	terminated: bool,
22	unreachable_hosts: HashMap<String, Host>,
23	}
24
25	const DEFAULT_FORKS: u32 = 5;
26
27	impl<'a> TaskQueueManager<'a> {
28	pub fn new(	×
29	forks: Option<u32>,
30	inventory_manager: &'a InventoryManager,
31	variable_manager: &'a VariableManager,
32	) -> Self {
33	Self {
34	callbacks: HashMap::new(),	×
35	callbacks_loaded: false,
36	forks: forks.unwrap_or(DEFAULT_FORKS),	×
37	inventory_manager,
38	variable_manager,
39	terminated: false,
40	unreachable_hosts: HashMap::new(),	×
41	}
42	}
43
44	pub fn inventory_manager(&self) -> &InventoryManager {	×
45	self.inventory_manager	×
46	}
47
48	pub fn variable_manager(&self) -> &VariableManager {	×
49	self.variable_manager	×
50	}
51
52	pub async fn run(&mut self, play: &Play) -> Result<()> {	×
53	self.load_callbacks(	×
54	// TODO: add logic to get callback plugin path
55	"/Users/darius/Programming/cogrs/dist/minimal-apple_x86_64-apple-darwin",
56	);
UNCOV 57	let all_vars = self	×
58	.variable_manager	×
59	.get_vars(Some(play), None, None, true, true);	×
60
61	self.emit_event(EventType::PlaybookOnPlayStart, None).await;	×
62
63	let mut play_iterator = PlayIterator::new(play);	×
64	play_iterator.init(self.inventory_manager)?;	×
65
66	let forks = min(self.forks, play_iterator.batch_size());	×
67
68	match play.strategy() {	×
UNCOV 69	Strategy::Linear => {	×
UNCOV 70	let mut strategy = LinearStrategy::new(&self);	×
UNCOV 71	strategy.run(&mut play_iterator)?;	×
72	}
73	Strategy::Free => {	×
74	todo!()
75	}
76	}
77
78	Ok(())	×
79	}
80
81	pub fn register_callback(&mut self, callback: Box<dyn CallbackPlugin>) {	×
82	let callback: Arc<dyn CallbackPlugin> = Arc::from(callback);	×
UNCOV 83	for event in callback.get_interested_events() {	×
UNCOV 84	self.callbacks	×
UNCOV 85	.entry(event)	×
86	.or_insert_with(Vec::new)	×
87	.push(callback.clone());	×
88	}
89	}
90
91	pub fn get_unreachable_hosts(&self) -> &HashMap<String, Host> {	×
UNCOV 92	&self.unreachable_hosts	×
93	}
94
UNCOV 95	pub fn is_terminated(&self) -> bool {	×
UNCOV 96	self.terminated	×
97	}
98
99	fn load_callbacks(&mut self, plugin_dir: &str) {	×
100	use libloading::{Library, Symbol};
101	use std::fs;
102
103	let plugin_extension = if cfg!(target_os = "windows") {
UNCOV 104	"dll"	×
105	} else if cfg!(target_os = "macos") {
106	"dylib"	×
107	} else {
UNCOV 108	"so"	×
109	};
110
111	if self.callbacks_loaded {	×
112	return;	×
113	}
114
UNCOV 115	for entry in fs::read_dir(plugin_dir).expect("Invalid plugin directory") {	×
UNCOV 116	let path = entry.expect("Failed to read entry").path();	×
117	if path.extension().and_then(\|e\| e.to_str()) == Some(plugin_extension) {	×
118	unsafe {
UNCOV 119	let lib = Library::new(&path).expect("Failed to load plugin");	×
120
121	// Dynamically load the callback creation function
UNCOV 122	let create_callback: Symbol<fn() -> Box<dyn CallbackPlugin>> = lib	×
123	.get(b"create_plugin")
124	.expect("Failed to find create_plugin function");
125
UNCOV 126	let plugin = create_callback();	×
127
128	// Register the plugin for events
129	self.register_callback(plugin);	×
130	}
131	}
132	}
133
UNCOV 134	self.callbacks_loaded = true	×
135	}
136
137	pub async fn emit_event(&self, event: EventType, data: Option<Value>) {	×
138	if let Some(callbacks) = self.callbacks.get(&event) {	×
139	// Spawn and collect tasks
140	let tasks: Vec<_> = callbacks	×
141	.iter()
142	.map(\|callback\| {	×
UNCOV 143	let callback = callback.clone();	×
UNCOV 144	let event = event.clone();	×
UNCOV 145	let data = data.clone();	×
UNCOV 146	tokio::spawn(async move {	×
UNCOV 147	callback.on_event(&event, data.as_ref()); // Async invocation	×
148	})
149	})
150	.collect();
151
152	// Wait for all spawned tasks to complete
UNCOV 153	for task in tasks {	×
UNCOV 154	if let Err(err) = task.await {	×
UNCOV 155	eprintln!("Callback task panicked: {:?}", err);	×
156	}
157	}
158	}
159	}
160	}

dariusbakunas / cogrs / 13381193139

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