manycastr/orchestrator/
worker.rs

1use crate::custom_module::manycastr::WorkerStatus::{Disconnected, Idle, Listening, Probing};
2use crate::custom_module::manycastr::{Address, ReplyBatch, WorkerStatus};
3use crate::orchestrator::{CliHandle, MeasurementHandle};
4use futures_core::Stream;
5use log::{debug, info, warn};
6use std::fmt;
7use std::pin::Pin;
8use std::sync::{Arc, Mutex};
9use std::task::{Context, Poll};
10use tokio::sync::mpsc;
11use tokio::sync::mpsc::Sender;
12
13/// Compare a `Mutex<WorkerStatus>` directly against a `WorkerStatus`
14impl PartialEq<WorkerStatus> for Mutex<WorkerStatus> {
15    fn eq(&self, other: &WorkerStatus) -> bool {
16        let status = self.lock().unwrap();
17        *status == *other
18    }
19}
20
21/// Special Receiver struct that notices when the worker disconnects.
22/// When a worker drops we update the active worker counter such that the orchestrator knows this worker is not participating in any measurements.
23/// Furthermore, we send a message to the CLI if it is currently performing a measurement, to let it know this worker is finished.
24pub struct WorkerReceiver<T> {
25    /// The inner receiver that connects to the worker
26    pub(crate) inner: mpsc::Receiver<T>,
27    /// All per-measurement state. `None` when idle.
28    pub(crate) measurement: MeasurementHandle,
29    /// Sender that connects to the CLI
30    pub(crate) cli_sender: CliHandle,
31    /// The hostname of the worker
32    pub(crate) hostname: String,
33    /// Worker ID
34    pub(crate) worker_id: u32,
35    /// The status of the worker, used to determine if it is connected or not
36    pub(crate) status: Arc<Mutex<WorkerStatus>>,
37}
38
39impl<T> Stream for WorkerReceiver<T> {
40    type Item = T;
41
42    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
43        self.inner.poll_recv(cx)
44    }
45}
46
47impl<T> Drop for WorkerReceiver<T> {
48    fn drop(&mut self) {
49        warn!("[Orchestrator] Worker {} lost connection", self.hostname);
50
51        let mut should_notify_cli = false;
52        let worker_id = self.worker_id;
53
54        {
55            let mut measurement_lock = self.measurement.write().unwrap();
56
57            // Do not wait for this disconnected worker for measurement finish
58            if let Some(ref mut state) = *measurement_lock
59                && let Some(participant) = state.participants.get_mut(&worker_id)
60                && participant.is_counted
61            {
62                participant.is_counted = false;
63
64                // Remove from probing list if they were a prober
65                state.probing_workers.retain(|&id| id != worker_id);
66
67                // Discard state owned by this worker so the measurement can still terminate
68                // TODO consider keeping follow-up tasks and trace sessions for reconnects
69                if let Some(stack) = state.worker_stacks.remove(&worker_id)
70                    && !stack.is_empty()
71                {
72                    warn!(
73                        "[Orchestrator] Discarding {} queued follow-up tasks for dropped worker {}",
74                        stack.len(),
75                        self.hostname
76                    );
77                }
78                if let Some(ref mut config) = state.trace_config {
79                    config
80                        .session_tracker
81                        .sessions
82                        .retain(|id, _| id.worker_id != worker_id);
83                }
84
85                if state.active_workers() == 0 {
86                    // This was the last worker still holding a completion claim
87                    info!(
88                        "[Orchestrator] Last active worker ({}) dropped. Measurement finished.",
89                        self.hostname
90                    );
91                    *measurement_lock = None; // Reset the state
92                    should_notify_cli = true;
93                }
94            }
95        }
96
97        // Set the status to Disconnected
98        *self.status.lock().unwrap() = Disconnected;
99
100        // Notify the CLI if the measurement is finished now
101        if should_notify_cli
102            && let Some(cli_tx_lock) = self.cli_sender.lock().unwrap().as_ref()
103            && let Err(e) = cli_tx_lock.try_send(Ok(ReplyBatch::default()))
104        {
105            warn!(
106                "[Orchestrator] Failed to send measurement finished signal to CLI: {}",
107                e
108            );
109        }
110    }
111}
112
113/// Special Sender struct for workers that sends tasks after a delay (based on the Worker interval).
114#[derive(Clone)]
115pub struct WorkerSender<T> {
116    /// Inner sender that connects to the orker
117    pub(crate) inner: Sender<T>,
118    /// Unique Worker ID
119    pub(crate) worker_id: u32,
120    /// Worker hostname
121    pub(crate) hostname: String,
122    /// Status of the Worker (e.g., Listening, Probing, Idle, Connected)
123    pub(crate) status: Arc<Mutex<WorkerStatus>>,
124    /// Unicast IPv4 address of the Worker (None if unavailable)
125    pub(crate) unicast_v4: Option<Address>,
126    /// Unicast IPv6 address of the Worker (None if unavailable)
127    pub(crate) unicast_v6: Option<Address>,
128}
129impl<T> WorkerSender<T> {
130    /// Checks if the sender is closed
131    pub fn is_closed(&self) -> bool {
132        self.inner.is_closed()
133    }
134
135    /// Sends an instruction to the worker.
136    /// On failure, logs a warning (once) and marks the worker as disconnected.
137    pub async fn send(&self, task: T) -> Result<(), mpsc::error::SendError<T>> {
138        match self.inner.send(task).await {
139            Ok(_) => Ok(()),
140            Err(e) => {
141                if self.get_status() == Disconnected {
142                    debug!(
143                        "[Orchestrator] Dropping send to disconnected worker {}",
144                        self.hostname
145                    );
146                } else {
147                    warn!(
148                        "[Orchestrator] Failed to send to worker {}: {e}",
149                        self.hostname
150                    );
151                    self.cleanup();
152                }
153                Err(e)
154            }
155        }
156    }
157
158    /// Sends an instruction to the worker without blocking
159    pub fn try_send(&self, task: T) -> Result<(), mpsc::error::TrySendError<T>> {
160        self.inner.try_send(task)
161    }
162
163    /// Marks the worker as disconnected
164    pub(crate) fn cleanup(&self) {
165        *self.status.lock().unwrap() = Disconnected;
166    }
167
168    pub fn is_participating(&self) -> bool {
169        let status = self.status.lock().unwrap();
170        *status == Probing || *status == Listening
171    }
172
173    pub fn get_status(&self) -> WorkerStatus {
174        *self.status.lock().unwrap()
175    }
176
177    /// The worker finished its measurement
178    pub fn finished(&self) {
179        let mut status = self.status.lock().unwrap();
180        // Set the status to Idle if it is not Disconnected
181        if *status != Disconnected {
182            *status = Idle;
183        }
184    }
185}
186impl<T> fmt::Debug for WorkerSender<T> {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        write!(
189            f,
190            "WorkerSender {{ worker_id: {}, hostname: {}, status: {} }}",
191            self.worker_id,
192            self.hostname,
193            self.get_status().as_str_name()
194        )
195    }
196}
manycastr/orchestrator/worker.rs

manycastr/orchestrator/
worker.rs
