studiole_command/services/
command_runner.rs

1#![allow(dead_code)]
2
3use crate::prelude::*;
4use tokio::sync::MutexGuard;
5
6/// Current state of the [`CommandRunner`].
7#[derive(Clone, Copy, Debug, Default, Eq, Error, PartialEq)]
8pub enum RunnerStatus {
9    #[default]
10    #[error("Runner is stopped")]
11    Stopped,
12    #[error("Stopping when the active commands are complete")]
13    Stopping,
14    #[error("Stopping when the queue is empty")]
15    Draining,
16    #[error("Running")]
17    Running,
18}
19
20/// Queue and execute commands across a pool of workers.
21pub struct CommandRunner<T: ICommandInfo> {
22    mediator: Arc<CommandMediator<T>>,
23    registry: Arc<CommandRegistry<T>>,
24    workers: Arc<WorkerPool<T>>,
25}
26
27impl<T: ICommandInfo + 'static> FromServicesAsync for CommandRunner<T> {
28    type Error = ResolveError;
29
30    async fn from_services_async(services: &ServiceProvider) -> Result<Self, Report<Self::Error>> {
31        Ok(Self::new(
32            services.get::<CommandMediator<T>>()?,
33            services.get_async::<CommandRegistry<T>>().await?,
34            services.get::<WorkerPool<T>>()?,
35        ))
36    }
37}
38
39impl<T: ICommandInfo + 'static> CommandRunner<T> {
40    /// Create a new [`CommandRunner`].
41    #[must_use]
42    pub fn new(
43        mediator: Arc<CommandMediator<T>>,
44        registry: Arc<CommandRegistry<T>>,
45        workers: Arc<WorkerPool<T>>,
46    ) -> Self {
47        Self {
48            mediator,
49            registry,
50            workers,
51        }
52    }
53
54    /// Start any number of workers.
55    ///
56    /// Each worker will have a unique ID.
57    ///
58    /// Status will be set to `Running`.
59    pub async fn start(&self, worker_count: usize) {
60        self.workers.start(worker_count).await;
61    }
62
63    /// Stop workers after draining the queue.
64    pub async fn drain(&self) {
65        self.mediator
66            .set_runner_status(RunnerStatus::Draining)
67            .await;
68        self.workers.wait_for_stop().await;
69    }
70
71    /// Stop workers after their current work is complete
72    pub async fn stop(&self) {
73        self.mediator
74            .set_runner_status(RunnerStatus::Stopping)
75            .await;
76        self.workers.wait_for_stop().await;
77    }
78
79    /// Queue a command as a request.
80    pub async fn queue_request<R: Executable + Into<T::Request> + Send + Sync + 'static>(
81        &self,
82        request: R,
83    ) -> Result<(), Report<QueueError>> {
84        trace!(%request, type = type_name::<R>(), "Queueing");
85        let command = self.registry.resolve(request.clone())?;
86        trace!(%request, type = type_name::<R>(), "Resolved command");
87        self.mediator.queue(request.into(), command).await;
88        Ok(())
89    }
90
91    /// Lock and return the current command status map.
92    ///
93    /// The [`MutexGuard`] must be dropped promptly or [`Worker`] execution will block.
94    pub async fn get_commands(&self) -> MutexGuard<'_, HashMap<T::Request, CommandStatus<T>>> {
95        self.mediator.get_commands().await
96    }
97
98    /// Take completed results for a specific request type.
99    ///
100    /// - Removes matching completed entries from the command map
101    /// - Entries still `Queued` or `Executing` are left in the map
102    pub async fn take_completed<R>(&self) -> Vec<(R, Result<R::Response, R::ExecutionError>)>
103    where
104        R: Executable + TryFrom<T::Request>,
105        R::Response: TryFrom<T::Success>,
106        R::ExecutionError: TryFrom<T::Failure>,
107    {
108        let mut commands = self.mediator.get_commands().await;
109        let keys: Vec<T::Request> = commands
110            .iter()
111            .filter(|(k, status)| {
112                R::try_from((*k).clone()).is_ok()
113                    && matches!(
114                        status,
115                        CommandStatus::Succeeded(_) | CommandStatus::Failed(_)
116                    )
117            })
118            .map(|(k, _)| k.clone())
119            .collect();
120        let mut results = Vec::with_capacity(keys.len());
121        for key in keys {
122            let Some(status) = commands.remove(&key) else {
123                unreachable!("already filtered to existing key");
124            };
125            let request = R::try_from(key)
126                .ok()
127                .expect("already filtered to matching variant");
128            let result = match status {
129                CommandStatus::Succeeded(success) => Ok(R::Response::try_from(success)
130                    .ok()
131                    .expect("request variant should match success variant")),
132                CommandStatus::Failed(failure) => Err(R::ExecutionError::try_from(failure)
133                    .ok()
134                    .expect("request variant should match failure variant")),
135                _ => unreachable!("filtered to completed only"),
136            };
137            results.push((request, result));
138        }
139        results
140    }
141
142    /// Take succeeded results for a specific request type.
143    ///
144    /// - Removes matching succeeded entries from the command map
145    /// - `Failed`, `Queued`, and `Executing` entries are left in the map
146    pub async fn take_succeeded<R>(&self) -> Vec<(R, R::Response)>
147    where
148        R: Executable + TryFrom<T::Request>,
149        R::Response: TryFrom<T::Success>,
150    {
151        let mut commands = self.mediator.get_commands().await;
152        let keys: Vec<T::Request> = commands
153            .iter()
154            .filter(|(k, status)| {
155                R::try_from((*k).clone()).is_ok() && matches!(status, CommandStatus::Succeeded(_))
156            })
157            .map(|(k, _)| k.clone())
158            .collect();
159        let mut results = Vec::with_capacity(keys.len());
160        for key in keys {
161            let Some(CommandStatus::Succeeded(success)) = commands.remove(&key) else {
162                unreachable!("already filtered to succeeded");
163            };
164            let request = R::try_from(key)
165                .ok()
166                .expect("already filtered to matching variant");
167            let response = R::Response::try_from(success)
168                .ok()
169                .expect("request variant should match success variant");
170            results.push((request, response));
171        }
172        results
173    }
174
175    /// Take failed results for a specific request type.
176    ///
177    /// - Removes matching failed entries from the command map
178    /// - `Succeeded`, `Queued`, and `Executing` entries are left in the map
179    pub async fn take_failed<R>(&self) -> Vec<(R, R::ExecutionError)>
180    where
181        R: Executable + TryFrom<T::Request>,
182        R::ExecutionError: TryFrom<T::Failure>,
183    {
184        let mut commands = self.mediator.get_commands().await;
185        let keys: Vec<T::Request> = commands
186            .iter()
187            .filter(|(k, status)| {
188                R::try_from((*k).clone()).is_ok() && matches!(status, CommandStatus::Failed(_))
189            })
190            .map(|(k, _)| k.clone())
191            .collect();
192        let mut results = Vec::with_capacity(keys.len());
193        for key in keys {
194            let Some(CommandStatus::Failed(failure)) = commands.remove(&key) else {
195                unreachable!("already filtered to failed");
196            };
197            let request = R::try_from(key)
198                .ok()
199                .expect("already filtered to matching variant");
200            let error = R::ExecutionError::try_from(failure)
201                .ok()
202                .expect("request variant should match failure variant");
203            results.push((request, error));
204        }
205        results
206    }
207}
208
209#[cfg(all(test, feature = "server"))]
210mod tests {
211    use super::*;
212
213    use std::time::Duration;
214    use tokio::time::sleep;
215
216    const WORKER_COUNT: usize = 3;
217    const A_COUNT: usize = 10;
218    const B_COUNT: usize = 10;
219    const A_DURATON: u64 = 100;
220    const B_DURATON: u64 = 100;
221    #[allow(clippy::as_conversions, clippy::integer_division)]
222    const A_TOTAL_DURATON: u64 = (A_COUNT / WORKER_COUNT) as u64 * A_DURATON;
223
224    #[tokio::test]
225    async fn command_runner() {
226        // Arrange
227        let services = ServiceBuilder::new().with_commands().build();
228        let runner = services
229            .get_async::<CommandRunner<CommandInfo>>()
230            .await
231            .expect("should be able to get runner");
232        let events = services
233            .get::<CommandEvents<CommandInfo>>()
234            .expect("should be able to get events");
235        events.start().await;
236        let _logger = init_test_logger();
237
238        // Act
239        runner.start(WORKER_COUNT).await;
240
241        info!("Adding {A_COUNT} commands to queue");
242        for i in 1..=A_COUNT {
243            let request = DelayRequest::new(format!("A{i}"), A_DURATON);
244            runner
245                .queue_request(request)
246                .await
247                .expect("should be able to queue command");
248        }
249        info!("Added {A_COUNT} commands to queue");
250
251        // Assert
252        let length = events
253            .count()
254            .await
255            .get_currently_queued()
256            .expect("should be able to subtract");
257        debug!("Queue: {length}");
258        // assert_eq!(length, A_COUNT, "Queue immediately after sending batch A");
259
260        wait(50).await;
261        let length = events
262            .count()
263            .await
264            .get_currently_queued()
265            .expect("should be able to subtract");
266        debug!("Queue: {length}");
267        assert_ne!(length, 0, "Queue soon after adding batch A");
268
269        wait(A_TOTAL_DURATON + 100).await;
270        let length = events
271            .count()
272            .await
273            .get_currently_queued()
274            .expect("should be able to subtract");
275        debug!("Queue: {length}");
276        assert_eq!(length, 0, "Queue after batch A should have completed");
277
278        info!("Adding {B_COUNT} commands to queue");
279        for i in 1..=B_COUNT {
280            let request = DelayRequest::new(format!("B{i}"), B_DURATON);
281            runner
282                .queue_request(request)
283                .await
284                .expect("should be able to queue command");
285        }
286        info!("Added {B_COUNT} commands to queue");
287
288        wait(50).await;
289        info!("Requesting stop");
290        runner.workers.stop().await;
291        info!("Completed stop");
292
293        let count = events.count().await;
294        let length = count
295            .get_currently_queued()
296            .expect("should be able to subtract");
297        debug!("Queue: {length}");
298        assert_eq!(length, 7, "Queue after stop");
299        let length = count.succeeded;
300        debug!("Succeeded: {length}");
301        assert_eq!(length, 13, "Succeeded after stop");
302    }
303
304    async fn wait(wait: u64) {
305        sleep(Duration::from_millis(wait)).await;
306        info!("Waiting {wait} ms");
307    }
308}
studiole_command/services/command_runner.rs

studiole_command/services/
command_runner.rs
