13441062843

Committed 20 Feb 2025 05:24PM UTC coverage: 36.493% (-1.0%) from 37.522%

Build # 13441062843

Build Type

push

github

Committed by

dariusbakunas

Commit Message

feat: add playcontext and start working on connection plugin

Coverage Stats

0 of 54 new or added lines in 5 files covered. (0.0%)

2 existing lines in 2 files now uncovered.

639 of 1751 relevant lines covered (36.49%)

1.22 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::playbook::play_context::{PlayContext, PlayContextBuilder};
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 {
    forks: usize,
    callbacks_loaded: bool,
    callbacks: HashMap<EventType, Vec<Arc<dyn CallbackPlugin>>>,
    terminated: bool,
    unreachable_hosts: HashMap<String, Host>,
    workers: Vec<tokio::task::JoinHandle<()>>,
    play_context: PlayContext,
}

const DEFAULT_FORKS: usize = 5;

impl TaskQueueManager {
    pub fn new(forks: Option<usize>, play_context: &PlayContext) -> 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)),
            play_context: (*play_context).clone(),
        }
    }

    pub fn play_context(&self) -> &PlayContext {
        &self.play_context
    }

    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().await?;
        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: Arc<dyn CallbackPlugin>) {
        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
    }

    async fn load_callbacks(&mut self) -> Result<()> {
        let plugin_loader = cogrs_plugins::plugin_loader::PluginLoader::instance();
        let loader = plugin_loader.lock().await;

        let plugins = loader.get_callback_plugins().await?;
        for plugin in plugins {
            self.register_callback(plugin);
        }

        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::playbook::play_context::{PlayContext, PlayContextBuilder};
6	use crate::strategy::linear::LinearStrategy;
7	use crate::strategy::Strategy;
8	use crate::vars::manager::VariableManager;
9	use anyhow::Result;
10	use cogrs_plugins::callback::{CallbackPlugin, EventType};
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	play_context: PlayContext,
24	}
25
26	const DEFAULT_FORKS: usize = 5;
27
28	impl TaskQueueManager {
NEW 29	pub fn new(forks: Option<usize>, play_context: &PlayContext) -> Self {	×
30	Self {
31	callbacks: HashMap::new(),	×
32	callbacks_loaded: false,
33	forks: forks.unwrap_or(DEFAULT_FORKS),	×
34	terminated: false,
35	unreachable_hosts: HashMap::new(),	×
36	workers: Vec::with_capacity(forks.unwrap_or(DEFAULT_FORKS)),	×
NEW 37	play_context: (*play_context).clone(),	×
38	}
39	}
40
NEW 41	pub fn play_context(&self) -> &PlayContext {	×
42	&self.play_context
43	}
44
45	pub fn get_worker(&mut self, index: usize) -> Option<&tokio::task::JoinHandle<()>> {	×
46	self.workers.get(index)	×
47	}
48
49	pub fn set_worker(&mut self, index: usize, worker: tokio::task::JoinHandle<()>) {	×
50	self.workers.insert(index, worker);	×
51	}
52
53	/// Iterates over the roles/tasks in a play, using the given (or default)
54	/// strategy for queueing tasks. The default is the linear strategy, which
55	/// operates like classic Ansible by keeping all hosts in lock-step with
56	/// a given task (meaning no hosts move on to the next task until all hosts
57	/// are done with the current task).
58	pub async fn run(	×
59	&mut self,
60	play: Play,
61	variable_manager: &VariableManager,
62	inventory_manager: &InventoryManager,
63	) -> Result<()> {
64	self.load_callbacks().await?;	×
65	let all_vars = variable_manager.get_vars(Some(&play), None, None, true, true);	×
66
67	self.emit_event(EventType::PlaybookOnPlayStart, None).await;	×
68
69	let strategy = *play.strategy();	×
70
71	let mut play_iterator = PlayIterator::new(play);	×
72	play_iterator.init(inventory_manager)?;	×
73
74	self.forks = min(self.forks, play_iterator.batch_size());	×
75
76	match strategy {	×
77	Strategy::Linear => {
78	let mut strategy = LinearStrategy::new(self, inventory_manager, variable_manager);	×
79	strategy.run(&mut play_iterator).await?;	×
80	}
81	Strategy::Free => {
82	todo!()
83	}
84	}
85
86	Ok(())	×
87	}
88
89	pub fn register_callback(&mut self, callback: Arc<dyn CallbackPlugin>) {	×
90	for event in callback.get_interested_events() {	×
91	self.callbacks	×
92	.entry(event)
93	.or_insert_with(Vec::new)
94	.push(callback.clone());	×
95	}
96	}
97
98	pub fn get_unreachable_hosts(&self) -> &HashMap<String, Host> {	×
99	&self.unreachable_hosts	×
100	}
101
102	pub fn is_terminated(&self) -> bool {	×
103	self.terminated	×
104	}
105
106	async fn load_callbacks(&mut self) -> Result<()> {	×
107	let plugin_loader = cogrs_plugins::plugin_loader::PluginLoader::instance();	×
108	let loader = plugin_loader.lock().await;	×
109
110	let plugins = loader.get_callback_plugins().await?;	×
111	for plugin in plugins {	×
112	self.register_callback(plugin);	×
113	}
114
115	self.callbacks_loaded = true;	×
116	Ok(())	×
117	}
118
119	pub async fn emit_event(&self, event: EventType, data: Option<Value>) {	×
120	if let Some(callbacks) = self.callbacks.get(&event) {	×
121	// Spawn and collect tasks
122	let tasks: Vec<_> = callbacks	×
123	.iter()
124	.map(\|callback\| {	×
125	let callback = callback.clone();	×
126	let event = event.clone();	×
127	let data = data.clone();	×
128	tokio::spawn(async move {	×
129	callback.on_event(&event, data.as_ref()); // Async invocation	×
130	})
131	})
132	.collect();
133
134	// Wait for all spawned tasks to complete
135	for task in tasks {	×
136	if let Err(err) = task.await {	×
137	eprintln!("Callback task panicked: {:?}", err);	×
138	}
139	}
140	}
141	}
142
143	pub fn forks(&self) -> usize {	×
144	self.forks	×
145	}
146	}

dariusbakunas / cogrs / 13441062843

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