manycastr/worker/
client.rs

1use crate::ALL_ORIGINS;
2use crate::custom_module;
3use crate::custom_module::manycastr::controller_client::ControllerClient;
4use crate::custom_module::manycastr::instruction::InstructionType;
5use crate::custom_module::manycastr::{Address, End, Start, Task, Tasks};
6use crate::tls::TlsOptions;
7use crate::worker::config::Worker;
8use local_ip_address::{local_ip, local_ipv6};
9use log::{info, warn};
10use std::error::Error;
11use std::sync::Arc;
12use std::sync::atomic::{AtomicBool, Ordering};
13use std::time::Duration;
14use tokio::sync::mpsc::Sender;
15use tonic::Request;
16use tonic::transport::Channel;
17
18/// Grace period after the last probe is sent before closing the listener
19const END_REPLY_GRACE_SECS: u64 = 1;
20
21impl Worker {
22    /// Connect to the orchestrator.
23    ///
24    /// # Arguments
25    /// * `address` - the address of the orchestrator in string format, containing both the address (IPv4, IPv6, or hostname) and port number
26    /// * `tls` - the TLS settings given on the command line
27    ///
28    /// # Returns
29    /// A gRPC client that is connected to the orchestrator
30    ///
31    /// # Remarks
32    /// When TLS is enabled, the connection is secured and the orchestrator authenticated.
33    pub(crate) async fn connect(
34        address: String,
35        tls: &TlsOptions<'_>,
36    ) -> Result<ControllerClient<Channel>, Box<dyn Error>> {
37        let scheme = if tls.is_enabled() { "https" } else { "http" };
38        let uri = format!("{scheme}://{address}");
39        let mut endpoint = Channel::from_shared(uri)?;
40
41        if tls.is_enabled() {
42            endpoint = endpoint.tls_config(tls.client_config(&address)?)?;
43        }
44
45        let channel = endpoint
46            .keep_alive_timeout(Duration::from_secs(30))
47            .http2_keep_alive_interval(Duration::from_secs(15))
48            .tcp_keepalive(Some(Duration::from_secs(60)))
49            .connect()
50            .await
51            .map_err(|e| format!("{e:?}"))?;
52
53        Ok(ControllerClient::new(channel))
54    }
55
56    /// Establish a formal connection with the orchestrator.
57    /// Obtains a unique worker ID from the orchestrator, establishes a stream for receiving tasks, and handles tasks as they come in.
58    pub(crate) async fn connect_to_server(&mut self) -> Result<(), Box<dyn Error>> {
59        let mut abort_outbound: Arc<AtomicBool> = Arc::new(AtomicBool::new(false)); // To force close outbound sending thread
60        let worker_req = custom_module::manycastr::Worker {
61            hostname: self.hostname.clone(),
62            worker_id: 0,
63            status: custom_module::manycastr::WorkerStatus::Idle as i32, // Placeholder status
64            unicast_v6: local_ipv6().ok().map(Address::from),
65            unicast_v4: local_ip().ok().map(Address::from),
66        };
67
68        // Establish stream of measurement instructions to the Orchestrator
69        let mut stream = self
70            .grpc_client
71            .worker_connect(Request::new(worker_req))
72            .await?
73            .into_inner();
74
75        // Obtain the unique worker ID set by the Orchestrator (first message)
76        let init_msg = stream
77            .message()
78            .await?
79            .ok_or("Stream closed before Init message received")?;
80
81        let worker_id = match init_msg.instruction_type {
82            Some(InstructionType::Init(init)) => init.worker_id as u16,
83            _ => return Err("Did not receive Init message from orchestrator".into()),
84        };
85        info!("[Worker] Connected to Orchestrator with assigned worker ID: {worker_id}");
86
87        // Await instructions
88        let mut probe_interval: u64 = 1;
89        while let Some(instruction) = stream.message().await? {
90            let instr_type = match instruction.instruction_type {
91                Some(it) => it,
92                None => {
93                    warn!("[Worker] Received empty instruction, skipping");
94                    continue;
95                }
96            };
97
98            // Check if we are currently busy with a measurement
99            let is_busy = self.is_busy.load(Ordering::SeqCst);
100
101            match (is_busy, instr_type) {
102                // Starting a measurement (whilst idle)
103                (false, InstructionType::Start(start)) => {
104                    abort_outbound = Arc::new(AtomicBool::new(false));
105                    probe_interval = start.probe_interval as u64;
106                    self.handle_start_instruction(start, worker_id, abort_outbound.clone())?;
107                }
108
109                // Ending a measurement (whilst busy)
110                (true, InstructionType::End(data)) => {
111                    self.handle_end_instruction(data, abort_outbound.clone())
112                        .await?;
113                }
114
115                // Receiving a new measurement (whilst busy) [INVALID]
116                (true, InstructionType::Start(_)) => {
117                    warn!("[Worker] Received new measurement while busy; ignoring.");
118                }
119
120                // Receiving a task batch (whilst busy): route tasks to the sender(s) of their origin
121                (true, InstructionType::Tasks(task_batch)) => {
122                    // Tasks with nprobes > 1 are re-sent every probe_interval seconds
123                    let mut repeats: Vec<Task> = task_batch
124                        .tasks
125                        .iter()
126                        .filter(|task| task.nprobes > 1)
127                        .copied()
128                        .collect();
129
130                    // Send the tasks to the appropriate origins
131                    route_tasks(&self.outbound_txs, task_batch).await;
132
133                    if !repeats.is_empty() {
134                        // Schedule the remaining sends for multi-probe tasks
135                        repeats.sort_unstable_by_key(|task| task.nprobes);
136                        schedule_repeats(self.outbound_txs.clone(), repeats, probe_interval);
137                    }
138                }
139
140                // Receiving any other instruction (whilst busy) [INVALID]
141                (true, _) => {
142                    warn!("[Worker] Received unexpected instruction while busy; ignoring.");
143                }
144
145                // Receiving anything but a new measurement (whilst idle) [INVALID]
146                (false, _) => {
147                    warn!("[Worker] Received task data while idle; ignoring.");
148                }
149            }
150        }
151        info!("[Worker] Stream closed by Orchestrator");
152        // TODO in-process reconnect
153
154        Ok(())
155    }
156
157    /// Start a new measurement.
158    /// Marks the worker as busy,
159    /// Initializes the abort signals to False (for outbound and inbound threads)
160    /// Calls the function to initialize the measurement
161    ///
162    /// # Arguments
163    /// `start` - The definition of the new measurement
164    /// `worker_id` - ID of this worker
165    /// `abort_outbound` - Abort signal to forcefully close the outbound thread
166    fn handle_start_instruction(
167        &mut self,
168        start: Start,
169        worker_id: u16,
170        abort_outbound: Arc<AtomicBool>,
171    ) -> Result<(), Box<dyn Error>> {
172        info!("[Worker] Starting measurement {}", start.m_id);
173
174        // Mark busy and reset the abort signal
175        self.is_busy.store(true, Ordering::SeqCst);
176        self.abort_inbound.store(false, Ordering::SeqCst);
177
178        // Initialize the measurement threads
179        self.init(start, worker_id, abort_outbound)?;
180        Ok(())
181    }
182
183    /// End an ongoing measurement.
184    ///
185    /// Graceful end (code 0): the outbound threads first drain any tasks still queued in their
186    /// channels, then the inbound listener stays open for a grace period to capture in-flight
187    /// replies before it is closed.
188    /// Forceful end (code != 0): outbound and inbound threads are closed immediately,
189    /// discarding any queued tasks.
190    ///
191    /// # Arguments
192    /// `end_instruction` - End instruction sent by the Orchestrator with an ending code
193    /// `abort_outbound` - Shared boolean to forcefully close the outbound/sending thread
194    async fn handle_end_instruction(
195        &mut self,
196        end_instruction: End,
197        abort_outbound: Arc<AtomicBool>,
198    ) -> Result<(), Box<dyn Error>> {
199        let is_graceful = end_instruction.code == 0;
200
201        if is_graceful {
202            info!("[Worker] Received finish signal");
203        } else {
204            warn!(
205                "[Worker] Received abort signal (code {})",
206                end_instruction.code
207            );
208            // Close inbound and outbound threads immediately (discard tasks left in the channel)
209            self.abort_inbound.store(true, Ordering::SeqCst);
210            abort_outbound.store(true, Ordering::SeqCst);
211        }
212
213        // Close outbound sending threads (gracefully); the End instruction is queued last
214        let txs = std::mem::take(&mut self.outbound_txs);
215        for (_, tx) in txs {
216            let _ = tx.send(InstructionType::End(end_instruction)).await;
217        }
218
219        let handles = std::mem::take(&mut self.outbound_handles);
220        if is_graceful {
221            // Close the listener only after all outbound threads have sent their tasks
222            let abort_inbound = self.abort_inbound.clone();
223            tokio::task::spawn_blocking(move || {
224                for handle in handles {
225                    let _ = handle.join();
226                }
227                std::thread::sleep(Duration::from_secs(END_REPLY_GRACE_SECS));
228                abort_inbound.store(true, Ordering::SeqCst);
229            });
230        }
231
232        Ok(())
233    }
234}
235
236/// Route a task batch to the outbound sender(s) of each task's origin.
237async fn route_tasks(outbound_txs: &[(u32, Sender<InstructionType>)], task_batch: Tasks) {
238    if let [(_, tx)] = outbound_txs {
239        // Simple forward when there is only a single origin
240        let _ = tx.send(InstructionType::Tasks(task_batch)).await;
241    } else {
242        // Forward when there is a matching origin_id attached, or ALL_ORIGINS is specified
243        for (origin_id, tx) in outbound_txs {
244            let tasks: Vec<Task> = task_batch
245                .tasks
246                .iter()
247                .filter(|t| t.origin_id == *origin_id || t.origin_id == ALL_ORIGINS)
248                .cloned()
249                .collect();
250            if !tasks.is_empty() {
251                let _ = tx.send(InstructionType::Tasks(Tasks { tasks })).await;
252            }
253        }
254    }
255}
256
257/// Re-send `repeats` every `probe_interval` seconds until every task has been
258/// sent `nprobes` times, spacing out the repeated probes.
259/// Stops early when the measurement ends (all outbound channels closed).
260fn schedule_repeats(
261    outbound_txs: Vec<(u32, Sender<InstructionType>)>,
262    mut repeats: Vec<Task>,
263    probe_interval: u64,
264) {
265    tokio::spawn(async move {
266        for round in 1u32.. {
267            tokio::time::sleep(Duration::from_secs(probe_interval)).await;
268            // Remove all finished multi-probe tasks (assumes a list sorted by nprobes)
269            repeats.drain(..repeats.partition_point(|task| task.nprobes <= round));
270            if repeats.is_empty() || outbound_txs.iter().all(|(_, tx)| tx.is_closed()) {
271                break; // All sends done, or the measurement ended
272            }
273            let tasks = repeats.clone();
274            route_tasks(&outbound_txs, Tasks { tasks }).await;
275        }
276    });
277}
manycastr/worker/client.rs

manycastr/worker/
client.rs
