Skip to main content

vv_agent/runtime/backends/
thread.rs

1use std::sync::Arc;
2use std::thread::{self, JoinHandle};
3
4use super::execute_cycle_loop;
5use crate::runtime::CancellationToken;
6use crate::types::{AgentResult, AgentTask, CycleRecord, Message, Metadata};
7
8#[derive(Debug, Clone)]
9pub struct ThreadBackend {
10    max_workers: usize,
11}
12
13impl Default for ThreadBackend {
14    fn default() -> Self {
15        Self::new(4)
16    }
17}
18
19impl ThreadBackend {
20    pub fn new(max_workers: usize) -> Self {
21        Self {
22            max_workers: max_workers.max(1),
23        }
24    }
25
26    pub fn max_workers(&self) -> usize {
27        self.max_workers
28    }
29
30    pub fn execute<F>(
31        &self,
32        _task: &AgentTask,
33        initial_messages: Vec<Message>,
34        shared_state: Metadata,
35        cycle_executor: F,
36        cancellation_token: Option<&CancellationToken>,
37        max_cycles: u32,
38    ) -> AgentResult
39    where
40        F: FnMut(
41            u32,
42            &mut Vec<Message>,
43            &mut Vec<CycleRecord>,
44            &mut Metadata,
45            Option<&CancellationToken>,
46        ) -> Option<AgentResult>,
47    {
48        execute_cycle_loop(
49            initial_messages,
50            shared_state,
51            cycle_executor,
52            cancellation_token,
53            max_cycles,
54        )
55    }
56
57    pub fn submit<R, F>(&self, function: F) -> JoinHandle<R>
58    where
59        R: Send + 'static,
60        F: FnOnce() -> R + Send + 'static,
61    {
62        thread::spawn(function)
63    }
64
65    pub fn parallel_map<T, R, F>(&self, function: F, items: Vec<T>) -> Vec<R>
66    where
67        T: Send + 'static,
68        R: Send + 'static,
69        F: Fn(T) -> R + Send + Sync + 'static,
70    {
71        let function = Arc::new(function);
72        let mut indexed_items = items.into_iter().enumerate();
73        let mut indexed_results = Vec::new();
74
75        loop {
76            let mut handles = Vec::new();
77            for _ in 0..self.max_workers {
78                let Some((index, item)) = indexed_items.next() else {
79                    break;
80                };
81                let function = Arc::clone(&function);
82                handles.push(thread::spawn(move || (index, function(item))));
83            }
84            if handles.is_empty() {
85                break;
86            }
87            for handle in handles {
88                indexed_results.push(handle.join().expect("thread backend worker panicked"));
89            }
90        }
91
92        indexed_results.sort_by_key(|(index, _)| *index);
93        indexed_results
94            .into_iter()
95            .map(|(_, result)| result)
96            .collect()
97    }
98}