manycastr/orchestrator/
mod.rs

1mod cli;
2mod config;
3mod result_handler;
4mod service;
5mod task_distributor;
6mod trace;
7mod worker;
8
9use std::collections::{HashMap, HashSet, VecDeque};
10use std::net::SocketAddr;
11use std::ops::AddAssign;
12use std::sync::{Arc, Mutex, RwLock};
13use std::time::{Duration, Instant};
14
15use crate::custom_module;
16use crate::custom_module::manycastr::{Address, MeasurementType, Start, WorkerStatus};
17use crate::orchestrator::config::{AllowedOrigin, load_allowed_origins, load_worker_config};
18use crate::orchestrator::mpsc::Sender;
19use crate::orchestrator::result_handler::SessionTracker;
20use crate::orchestrator::worker::WorkerSender;
21use crate::tls::server_identity;
22use clap::ArgMatches;
23use custom_module::manycastr::{
24    Instruction, ReplyBatch, Task, controller_server::ControllerServer,
25};
26use log::{info, warn};
27use tokio::sync::mpsc;
28use tonic::codec::CompressionEncoding;
29use tonic::transport::{Identity, ServerTlsConfig};
30use tonic::{Status, transport::Server};
31
32type ResultMessage = Result<ReplyBatch, Status>;
33type CliSender = Sender<ResultMessage>;
34pub(crate) type CliHandle = Arc<Mutex<Option<CliSender>>>;
35
36type TaskMessage = Result<Instruction, Status>;
37
38/// Shared registry of connected worker senders. Updated when a worker reconnects.
39pub(crate) type WorkerRegistry = Arc<Mutex<Vec<WorkerSender<TaskMessage>>>>;
40
41/// Shared handle to the active measurement state. `None` when no measurement is running.
42pub type MeasurementHandle = Arc<RwLock<Option<MeasurementState>>>;
43
44/// A worker participating in the active measurement.
45/// Participants can re-join after disconnect.
46///
47/// The entry is removed when the worker finishes.
48#[derive(Debug)]
49pub struct Participant {
50    /// The worker's role in the measurement (Probing or Listening)
51    pub role: WorkerStatus,
52    /// When true, the Orchestrator waits for this participant before measurement finish.
53    pub is_counted: bool,
54}
55
56/// All state associated with a single active measurement.
57#[derive(Debug)]
58pub struct MeasurementState {
59    /// 16-bit measurement ID (filtering replies) + 16-bit session ID (reply attribution).
60    pub m_id: u32,
61    /// Worker IDs of connected Workers that are actively probing
62    pub probing_workers: Vec<u32>,
63    /// Participating workers (removed when a worker finishes; kept on disconnect for rejoin)
64    pub participants: HashMap<u32, Participant>,
65    /// Per-worker Start instructions (re-sent when a worker rejoins mid-measurement)
66    pub start_instructions: HashMap<u32, Start>,
67    /// Whether the current measurement is being finalized (no new tasks being sent)
68    pub is_finalizing: bool,
69    /// The measurement type (LACeS, catchment, latency, …)
70    pub m_type: MeasurementType,
71    /// Whether `--responsive` gates a broadcast of measurement probes.
72    pub is_gated_broadcast: bool,
73    /// Whether the hitlist holds ranked candidates per prefix (ISI format)
74    pub is_prefix_hitlist: bool,
75    /// Number of times each measurement probe is sent (always >= 1)
76    pub nprobes: u32,
77    /// Per-worker stacks of follow-up tasks (discovery → measurement, traceroute hops)
78    pub worker_stacks: HashMap<u32, VecDeque<Task>>,
79    /// Traceroute configuration and session tracker (None for non-traceroute measurements)
80    pub trace_config: Option<TracerouteConfig>,
81    /// Resolved --responsive targets
82    pub resolved_targets: HashSet<Address>,
83    /// Live feed state (None for hitlist-based measurements)
84    pub live: Option<LiveState>,
85}
86
87impl MeasurementState {
88    /// Number of connected workers participating in a measurement.
89    /// The measurement is complete when this reaches zero.
90    pub fn active_workers(&self) -> usize {
91        self.participants.values().filter(|p| p.is_counted).count()
92    }
93}
94
95/// Encode nprobes 1 as 0 for gRPC compression
96#[inline]
97pub fn wire_nprobes(nprobes: u32) -> u32 {
98    if nprobes > 1 { nprobes } else { 0 }
99}
100
101/// Timeout for live-feed discovery probes
102pub const LIVE_DISCOVERY_TIMEOUT_SECS: u64 = 3;
103
104/// Worker selection of a live-feed target (parsed from `LiveTarget.worker_ids`).
105#[derive(Debug)]
106pub enum WorkerSel {
107    /// Any single worker (round-robin over probing workers)
108    Any,
109    /// All probing workers (staggered broadcast)
110    All,
111    /// An explicit set of workers, staggered like a broadcast (sorted and deduplicated)
112    Set(Vec<u32>),
113}
114
115impl WorkerSel {
116    /// Whether the selection targets more than one worker.
117    pub fn is_multi(&self) -> bool {
118        match self {
119            WorkerSel::Any => false,
120            WorkerSel::All => true,
121            WorkerSel::Set(ids) => ids.len() > 1,
122        }
123    }
124}
125
126/// State for a live (feed-based) measurement.
127#[derive(Debug)]
128pub struct LiveState {
129    /// Pending discovery probes awaiting a response tracked by address and session ID.
130    pub pending: HashMap<(Address, u32), PendingTarget>,
131    /// Follow-up task stacks for explicit worker sets (sent staggered like a broadcast).
132    pub set_stacks: HashMap<Vec<u32>, VecDeque<Task>>,
133    /// Recently dispatched trace targets and their reply deadline (feed-trace only).
134    pub trace_targets: HashMap<Address, Instant>,
135}
136
137impl LiveState {
138    /// When receiving a --discovery probe reply, remove the associated pending target.
139    /// Makes use of session_id when used in discovery probes (ICMP/DNS-A only).
140    pub fn remove_pending(
141        &mut self,
142        addr: Address,
143        session_id: u32,
144    ) -> Option<(u32, PendingTarget)> {
145        if let Some(pending) = self.pending.remove(&(addr, session_id)) {
146            return Some((session_id, pending));
147        }
148        if session_id == 0 {
149            let key = self.pending.keys().find(|(a, _)| *a == addr).copied()?;
150            let pending = self.pending.remove(&key)?;
151            return Some((key.1, pending));
152        }
153        None
154    }
155}
156
157/// A live target awaiting a discovery reply before it is probed (or given up on timeout).
158#[derive(Debug)]
159pub struct PendingTarget {
160    /// Worker selection for the follow-up measurement probes
161    pub worker_sel: WorkerSel,
162    /// Worker performing the discovery probe
163    pub discovery_worker: u32,
164    /// Number of measurement probes to send (per worker) once the target resolves (always >= 1)
165    pub nprobes: u32,
166    /// When the discovery attempt expires
167    pub deadline: Instant,
168}
169
170/// Traceroute configuration
171#[derive(Debug)]
172pub struct TracerouteConfig {
173    /// Session tracker for Trace Tasks
174    pub session_tracker: SessionTracker,
175    /// Origin tracemap seed probes are sent from (anycast-traceroute sessions
176    /// instead use the origin that caught the discovery reply)
177    pub origin_id: u32,
178    /// Timeout value for traceroute measurements (default 3s)
179    pub timeout: u64,
180    /// Max hop count for traceroute measurements (default 25)
181    pub max_hops: u32,
182    /// Hop count to start traceroute measurements with (default 4)
183    pub initial_hop: u32,
184    /// Maximum number of unresponsive hops before terminating the traceroute
185    /// (default 5; tracemap confirmation window: 3)
186    pub max_failures: u32,
187    /// Whether to emit a '*' hop (no reply) to the CLI when a hop times out
188    pub star_unresponsive: bool,
189}
190
191/// Which client(s) this Orchestrator service accepts.
192///
193/// Default: a single service that accepts both Workers and the CLI.
194/// With `--cli_port` create a separate Orchestrator service for Workers and CLI.
195#[derive(Debug, Clone, Copy, PartialEq, Eq)]
196pub enum Access {
197    Both,
198    Workers,
199    Cli,
200}
201
202impl Access {
203    /// Accept Worker calls (`worker_connect`, `send_result`, `measurement_finished`)
204    fn serves_workers(self) -> bool {
205        matches!(self, Access::Both | Access::Workers)
206    }
207
208    /// Accept CLI calls (`do_measurement`, `live_measurement`, `list_workers`)
209    fn serves_cli(self) -> bool {
210        matches!(self, Access::Both | Access::Cli)
211    }
212}
213
214/// The main orchestrator service struct.
215#[derive(Debug, Clone)]
216pub struct ControllerService {
217    /// List of connected workers
218    saved_workers: WorkerRegistry,
219    /// Sender to the CLI for streaming results
220    cli_sender: CliHandle,
221    /// All per-measurement state. `None` when idle.
222    measurement: MeasurementHandle,
223    /// Last used unique worker ID
224    unique_id: Arc<Mutex<u32>>,
225    /// Optional static mapping of hostnames to worker IDs
226    worker_config: Option<HashMap<String, u32>>,
227    /// Maximum probing rate (probes per second, per worker) allowed for measurements (None = unlimited)
228    max_rate: Option<u32>,
229    /// Optional allow-list of origins CLIs may use (None = all origins allowed)
230    allowed_origins: Option<Vec<AllowedOrigin>>,
231    /// Which clients this instance serves
232    access: Access,
233}
234
235impl ControllerService {
236    /// Refuse Worker calls on a CLI-only Orc service
237    fn check_worker_access(&self, call: &str) -> Result<(), Status> {
238        if self.access.serves_workers() {
239            return Ok(());
240        }
241
242        warn!("[Orchestrator] Refused worker call '{call}' on the CLI port");
243        Err(Status::permission_denied(
244            "this port only serves the CLI, connect workers to the worker port",
245        ))
246    }
247
248    /// Refuse CLI calls on a Worker-only Orc service
249    fn check_cli_access(&self, call: &str) -> Result<(), Status> {
250        if self.access.serves_cli() {
251            return Ok(());
252        }
253
254        warn!("[Orchestrator] Refused CLI call '{call}' on the worker port");
255        Err(Status::permission_denied(
256            "this port only serves workers, connect the CLI to the CLI port",
257        ))
258    }
259
260    /// Gets a unique worker ID for a new connecting worker.
261    /// Increments the unique ID counter after returning the ID (for the next worker).
262    fn get_unique_id(&self) -> u32 {
263        let mut unique_id = self.unique_id.lock().unwrap();
264        let worker_id = *unique_id;
265        unique_id.add_assign(1);
266
267        worker_id
268    }
269
270    /// Gets a worker ID for a connecting worker based on its hostname.
271    /// If the hostname already exists, it returns the existing worker ID.
272    /// If the hostname does not exist, it checks for a statically configured ID or generates a new unique ID.
273    ///
274    /// # Arguments
275    /// * `hostname` - the hostname of the worker
276    ///
277    /// # Returns
278    /// A tuple containing: the worker ID and a boolean indicating if this is a reconnection of a closed worker.
279    ///
280    /// # Errors
281    /// Returns an error if the hostname already exists and is used by a connected worker.
282    fn get_worker_id(&self, hostname: &str) -> Result<(u32, bool), Status> {
283        {
284            let workers = self.saved_workers.lock().unwrap();
285            // Check if the hostname already exists in the workers list
286            if let Some(existing_worker) = workers.iter().find(|w| w.hostname == hostname) {
287                return if !existing_worker.is_closed() {
288                    warn!("[Orchestrator] Refusing worker, hostname already exists: {hostname}");
289                    Err(Status::already_exists("This hostname already exists"))
290                } else {
291                    // This is a reconnection of a closed worker.
292                    let id = existing_worker.worker_id;
293                    Ok((id, true))
294                };
295            }
296        }
297
298        // Check for a statically configured ID
299        if let Some(worker_config) = &self.worker_config
300            && let Some(worker_id) = worker_config.get(hostname)
301        {
302            return Ok((*worker_id, false));
303        }
304
305        // Return a new unique ID
306        let new_id = self.get_unique_id();
307        Ok((new_id, false))
308    }
309
310    /// Get a random measurement ID (u16)
311    fn next_m_id(&self) -> u32 {
312        rand::random::<u16>() as u32
313    }
314}
315
316/// Starts the orchestrator on the specified port.
317///
318/// Workers and CLIs share one port by default.
319/// With `--cli_port` the orchestrator listens on two ports instead.
320///
321/// # Arguments
322/// * `args` - the parsed command-line arguments
323///
324/// # Errors
325/// If a listening address cannot be parsed, or a listener cannot be served.
326pub async fn start(args: &ArgMatches) -> Result<(), Box<dyn std::error::Error>> {
327    let port = *args.get_one::<u16>("port").unwrap();
328    let cli_port = args.get_one::<u16>("cli_port").copied();
329    if cli_port == Some(port) {
330        return Err(
331            format!("--cli_port {port} must differ from --port, or be left out entirely").into(),
332        );
333    }
334
335    // Get optional configuration file
336    let (current_worker_id, worker_config) = args
337        .get_one::<String>("config")
338        .map(load_worker_config)
339        .unwrap_or_else(|| (Arc::new(Mutex::new(1)), None));
340
341    let controller = ControllerService {
342        saved_workers: Arc::new(Mutex::new(Vec::new())),
343        cli_sender: Arc::new(Mutex::new(None)),
344        measurement: Arc::new(RwLock::new(None)),
345        unique_id: current_worker_id,
346        worker_config,
347        max_rate: args.get_one::<u32>("max_rate").copied(),
348        allowed_origins: args.get_one::<String>("origins").map(load_allowed_origins),
349        access: Access::Both,
350    };
351
352    // if TLS is enabled create the orchestrator using a TLS configuration
353    let identity = args.get_one::<String>("tls").map(|cert_path| {
354        info!("[Orchestrator] Starting orchestrator with TLS enabled");
355        server_identity(
356            cert_path,
357            args.get_one::<String>("tls_key").map(String::as_str),
358        )
359    });
360
361    let Some(cli_port) = cli_port else {
362        info!("[Orchestrator] Serving Workers and CLIs on port {port}");
363        return serve(port, controller, identity).await;
364    };
365
366    // Separate ports for Workers and CLI clients
367    info!("[Orchestrator] Serving Workers on port {port}, CLIs on port {cli_port}");
368    let workers = serve(
369        port,
370        ControllerService {
371            access: Access::Workers,
372            ..controller.clone()
373        },
374        identity.clone(),
375    );
376    let clis = serve(
377        cli_port,
378        ControllerService {
379            access: Access::Cli,
380            ..controller
381        },
382        identity,
383    );
384
385    // Either listener failing takes the orchestrator down
386    tokio::try_join!(workers, clis)?;
387
388    Ok(())
389}
390
391/// Serve the controller on `[::]:port`, optionally with TLS.
392///
393/// # Arguments
394/// * `port` - the port to listen on
395/// * `controller` - the Orchestrator service instance
396/// * `identity` - the TLS certificate and private key, or `None` when using no TLS
397///
398/// # Errors
399/// If the address cannot be parsed, the TLS identity is rejected, or the port cannot be served.
400async fn serve(
401    port: u16,
402    controller: ControllerService,
403    identity: Option<Identity>,
404) -> Result<(), Box<dyn std::error::Error>> {
405    let addr: SocketAddr = format!("[::]:{port}").parse()?;
406
407    let svc = ControllerServer::new(controller)
408        .accept_compressed(CompressionEncoding::Zstd)
409        .max_decoding_message_size(10 * 1024 * 1024 * 1024) // 10 GB
410        .max_encoding_message_size(10 * 1024 * 1024 * 1024);
411
412    let mut builder = Server::builder();
413    if let Some(identity) = identity {
414        builder = builder.tls_config(ServerTlsConfig::new().identity(identity))?;
415    }
416
417    builder
418        .http2_keepalive_interval(Some(Duration::from_secs(10)))
419        .http2_keepalive_timeout(Some(Duration::from_secs(20)))
420        .tcp_keepalive(Some(Duration::from_secs(30)))
421        .add_service(svc)
422        .serve(addr)
423        .await
424        .map_err(|e| format!("unable to serve on port {port}: {}", &e))?;
425
426    Ok(())
427}
manycastr/orchestrator/mod.rs

manycastr/orchestrator/
mod.rs
