manycastr/orchestrator/
cli.rs

1use crate::custom_module::manycastr::WorkerStatus::Disconnected;
2use crate::custom_module::manycastr::{End, Instruction, instruction};
3use crate::orchestrator::{MeasurementHandle, WorkerRegistry};
4use futures_core::Stream;
5use log::warn;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8use tokio::sync::mpsc;
9
10/// Special Receiver struct that notices when the CLI disconnects.
11/// When a CLI disconnects we cancel the measurement it was performing (if still active):
12/// the measurement state is cleared and all participating workers are sent an abort
13/// instruction, making the orchestrator and workers available for a new measurement.
14pub struct CLIReceiver<T> {
15    /// Receiver that connects to the CLI
16    pub(crate) inner: mpsc::Receiver<T>,
17    /// All per-measurement state. `None` when idle.
18    pub(crate) measurement: MeasurementHandle,
19    /// Registry of connected workers (for aborting the measurement on the workers)
20    pub(crate) workers: WorkerRegistry,
21    /// ID of the measurement this CLI stream belongs to
22    pub(crate) m_id: u32,
23}
24
25impl<T> Stream for CLIReceiver<T> {
26    type Item = T;
27
28    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
29        self.inner.poll_recv(cx)
30    }
31}
32
33impl<T> Drop for CLIReceiver<T> {
34    fn drop(&mut self) {
35        // Clear the measurement state, but only if our measurement is still the active one
36        let participant_ids: Vec<u32> = {
37            let mut lock = self.measurement.write().unwrap();
38            match lock.as_ref() {
39                Some(state) if state.m_id == self.m_id => {
40                    warn!(
41                        "[Orchestrator] CLI dropped during an active measurement, terminating measurement"
42                    );
43                    let ids = state.participants.keys().copied().collect();
44                    *lock = None; // No longer an active measurement
45                    ids
46                }
47                // Our measurement already finished (or was replaced by a new one)
48                _ => return,
49            }
50        };
51
52        // Abort the measurement on all participating workers
53        let abort = Instruction {
54            instruction_type: Some(instruction::InstructionType::End(End { code: 1 })),
55        };
56        let senders: Vec<_> = self.workers.lock().unwrap().clone();
57        for sender in senders {
58            if !participant_ids.contains(&sender.worker_id) || sender.get_status() == Disconnected {
59                continue;
60            }
61            if sender.try_send(Ok(abort.clone())).is_err() {
62                warn!(
63                    "[Orchestrator] Could not send abort instruction to worker {}",
64                    sender.hostname
65                );
66            }
67            sender.finished();
68        }
69    }
70}
manycastr/orchestrator/cli.rs

manycastr/orchestrator/
cli.rs
