laminar_core/cluster/discovery/gossip_discovery/
mod.rs

1//! Gossip-based discovery using chitchat.
2//!
3//! Uses the chitchat protocol (from Quickwit) for decentralized
4//! node discovery with phi-accrual failure detection.
5
6#![allow(clippy::disallowed_types)] // cold path: gossip discovery coordination
7use std::collections::{BTreeMap, HashMap};
8#[cfg(feature = "cluster")]
9use std::net::ToSocketAddrs;
10use std::sync::Arc;
11use std::time::Duration;
12
13use parking_lot::RwLock;
14use tokio::sync::watch;
15use tokio_util::sync::CancellationToken;
16
17use super::{Discovery, DiscoveryError, NodeId, NodeInfo, NodeMetadata, NodeState};
18
19const MAX_METADATA_TAGS: usize = 32;
20const MAX_METADATA_TAG_KEY_BYTES: usize = 128;
21const MAX_METADATA_TAG_VALUE_BYTES: usize = 1_024;
22const MAX_METADATA_TAGS_ENCODED_BYTES: usize = 8 * 1_024;
23const PROCESS_INCARNATION_TAG: &str = "laminardb.process-incarnation";
24const DISCOVERY_PROTOCOL_VERSION: &str = "2";
25const DISCOVERY_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
26
27struct AbortTaskOnDrop(tokio::task::AbortHandle);
28
29impl AbortTaskOnDrop {
30    fn abort(&self) {
31        self.0.abort();
32    }
33}
34
35impl Drop for AbortTaskOnDrop {
36    fn drop(&mut self) {
37        self.abort();
38    }
39}
40
41async fn join_task_bounded<T>(
42    mut task: tokio::task::JoinHandle<T>,
43    timeout: Duration,
44    task_name: &'static str,
45) -> Option<Result<T, tokio::task::JoinError>> {
46    let abort_on_drop = AbortTaskOnDrop(task.abort_handle());
47    if let Ok(result) = tokio::time::timeout(timeout, &mut task).await {
48        Some(result)
49    } else {
50        tracing::warn!(
51            task = task_name,
52            ?timeout,
53            "Discovery task did not stop in time"
54        );
55        abort_on_drop.abort();
56        let _ = tokio::time::timeout(timeout.min(Duration::from_secs(1)), &mut task).await;
57        None
58    }
59}
60
61struct ChitchatShutdownGuard(chitchat::ChitchatHandle);
62
63impl ChitchatShutdownGuard {
64    fn new(handle: chitchat::ChitchatHandle) -> Self {
65        Self(handle)
66    }
67
68    fn handle(&self) -> &chitchat::ChitchatHandle {
69        &self.0
70    }
71}
72
73impl Drop for ChitchatShutdownGuard {
74    fn drop(&mut self) {
75        self.0.abort();
76    }
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80struct HighestPeerGeneration {
81    generation: u64,
82    ambiguous: bool,
83}
84
85fn observe_peer_generation(
86    peers: &mut HashMap<u64, HighestPeerGeneration>,
87    node_id: u64,
88    generation: u64,
89) {
90    use std::collections::hash_map::Entry;
91
92    match peers.entry(node_id) {
93        Entry::Vacant(entry) => {
94            entry.insert(HighestPeerGeneration {
95                generation,
96                ambiguous: false,
97            });
98        }
99        Entry::Occupied(mut entry) => {
100            let current = *entry.get();
101            if generation > current.generation {
102                entry.insert(HighestPeerGeneration {
103                    generation,
104                    ambiguous: false,
105                });
106            } else if generation == current.generation {
107                tracing::warn!(
108                    node_id,
109                    generation,
110                    "ambiguous gossip process identity at one generation; excluding stable node"
111                );
112                entry.insert(HighestPeerGeneration {
113                    generation,
114                    ambiguous: true,
115                });
116            }
117        }
118    }
119}
120
121fn stable_node_id(node_id: &str) -> Option<u64> {
122    node_id.strip_prefix("node-")?.parse().ok()
123}
124
125fn is_node_info_key(key: &str) -> bool {
126    matches!(
127        key,
128        keys::NODE_STATE
129            | keys::RPC_ADDRESS
130            | keys::NODE_NAME
131            | keys::LOAD_CORES
132            | keys::LOAD_MEMORY
133            | keys::FAILURE_DOMAIN
134            | keys::NODE_VERSION
135            | keys::PROTOCOL_VERSION
136            | keys::METADATA_TAGS
137    )
138}
139
140/// Key namespace for chitchat key-value pairs.
141pub mod keys {
142    /// Node state key.
143    pub const NODE_STATE: &str = "node:state";
144    /// RPC address key.
145    pub const RPC_ADDRESS: &str = "node:rpc_addr";
146    /// Node name key.
147    pub const NODE_NAME: &str = "node:name";
148    /// CPU core count key.
149    pub const LOAD_CORES: &str = "load:cores";
150    /// Memory bytes key.
151    pub const LOAD_MEMORY: &str = "load:memory_bytes";
152    /// Failure domain key.
153    pub const FAILURE_DOMAIN: &str = "node:failure_domain";
154    /// Version key.
155    pub const NODE_VERSION: &str = "node:version";
156    /// Exact discovery wire revision.
157    pub const PROTOCOL_VERSION: &str = "node:discovery_protocol";
158    /// Canonical JSON object containing the complete user/runtime metadata-tag map.
159    pub const METADATA_TAGS: &str = "node:metadata_tags";
160}
161
162/// Configuration for gossip-based discovery.
163#[derive(Debug, Clone)]
164pub struct GossipDiscoveryConfig {
165    /// Address to bind the gossip listener.
166    pub gossip_address: String,
167    /// Seed node addresses for initial cluster bootstrap.
168    pub seed_nodes: Vec<String>,
169    /// Interval between gossip rounds.
170    pub gossip_interval: Duration,
171    /// Phi-accrual failure detector threshold.
172    pub phi_threshold: f64,
173    /// Grace period before removing dead nodes.
174    pub dead_node_grace_period: Duration,
175    /// Cluster identifier (must match across all nodes).
176    pub cluster_id: String,
177    /// This node's ID.
178    pub node_id: NodeId,
179    /// Durable, monotonically increasing process term for this stable node ID.
180    pub process_generation: u64,
181    /// This node's info (published via chitchat keys).
182    pub local_node: NodeInfo,
183    /// Optional hostname or IP to advertise.
184    pub advertise_host: Option<String>,
185}
186
187impl Default for GossipDiscoveryConfig {
188    fn default() -> Self {
189        let mut metadata = NodeMetadata {
190            version: env!("CARGO_PKG_VERSION").into(),
191            ..NodeMetadata::default()
192        };
193        metadata.tags.insert(
194            PROCESS_INCARNATION_TAG.into(),
195            uuid::Uuid::from_u128(1).to_string(),
196        );
197        Self {
198            gossip_address: "127.0.0.1:9003".into(),
199            seed_nodes: Vec::new(),
200            gossip_interval: Duration::from_millis(500),
201            phi_threshold: 8.0,
202            dead_node_grace_period: Duration::from_secs(300),
203            cluster_id: "laminardb-default".into(),
204            node_id: NodeId(1),
205            process_generation: 1,
206            local_node: NodeInfo {
207                id: NodeId(1),
208                name: "node-1".into(),
209                rpc_address: "127.0.0.1:9000".into(),
210                state: NodeState::Active,
211                metadata,
212                last_heartbeat_ms: 0,
213            },
214            advertise_host: None,
215        }
216    }
217}
218
219/// Gossip-based discovery using the chitchat protocol.
220pub struct GossipDiscovery {
221    config: GossipDiscoveryConfig,
222    peers: Arc<RwLock<HashMap<u64, NodeInfo>>>,
223    membership_tx: watch::Sender<Vec<NodeInfo>>,
224    membership_rx: watch::Receiver<Vec<NodeInfo>>,
225    cancel: CancellationToken,
226    started: bool,
227    chitchat_handle: Option<chitchat::ChitchatHandle>,
228    membership_handle: Option<tokio::task::JoinHandle<()>>,
229}
230
231impl GossipDiscovery {
232    /// Create a new gossip discovery instance.
233    #[must_use]
234    pub fn new(config: GossipDiscoveryConfig) -> Self {
235        let (tx, rx) = watch::channel(Vec::new());
236        Self {
237            config,
238            peers: Arc::new(RwLock::new(HashMap::new())),
239            membership_tx: tx,
240            membership_rx: rx,
241            cancel: CancellationToken::new(),
242            started: false,
243            chitchat_handle: None,
244            membership_handle: None,
245        }
246    }
247
248    /// Borrow the underlying chitchat handle, if the discovery has
249    /// been started. Enables other cluster components (barrier
250    /// coordinator, shuffle peer registry) to share the same chitchat
251    /// instance rather than spawning their own.
252    #[must_use]
253    pub fn chitchat_handle(&self) -> Option<&chitchat::ChitchatHandle> {
254        self.chitchat_handle.as_ref()
255    }
256
257    /// Parse a `NodeInfo` from chitchat key-value pairs.
258    fn parse_node_info(node_id_str: &str, kvs: &HashMap<String, String>) -> Option<NodeInfo> {
259        let id = stable_node_id(node_id_str)?;
260        let rpc_address = kvs.get(keys::RPC_ADDRESS)?.clone();
261        let name = kvs
262            .get(keys::NODE_NAME)
263            .filter(|name| !name.is_empty())?
264            .clone();
265        let state = kvs.get(keys::NODE_STATE).and_then(|s| match s.as_str() {
266            "joining" => Some(NodeState::Joining),
267            "active" => Some(NodeState::Active),
268            "suspected" => Some(NodeState::Suspected),
269            "draining" => Some(NodeState::Draining),
270            "left" => Some(NodeState::Left),
271            _ => None,
272        })?;
273
274        let cores: u32 = kvs.get(keys::LOAD_CORES)?.parse().ok()?;
275        let memory_bytes: u64 = kvs.get(keys::LOAD_MEMORY)?.parse().ok()?;
276        let failure_domain = kvs.get(keys::FAILURE_DOMAIN).cloned();
277        let version = kvs
278            .get(keys::NODE_VERSION)
279            .filter(|version| version.as_str() == env!("CARGO_PKG_VERSION"))?
280            .clone();
281        kvs.get(keys::PROTOCOL_VERSION)
282            .filter(|version| version.as_str() == DISCOVERY_PROTOCOL_VERSION)?;
283        let tags: HashMap<String, String> = {
284            let encoded = kvs.get(keys::METADATA_TAGS)?;
285            if encoded.len() > MAX_METADATA_TAGS_ENCODED_BYTES {
286                return None;
287            }
288            let tags = serde_json::from_str(encoded).ok()?;
289            Self::validate_metadata_tags(&tags, encoded.len()).ok()?;
290            tags
291        };
292        Self::validate_process_incarnation(&tags).ok()?;
293
294        Some(NodeInfo {
295            id: NodeId(id),
296            name,
297            rpc_address,
298            state,
299            metadata: NodeMetadata {
300                cores,
301                memory_bytes,
302                failure_domain,
303                tags,
304                version,
305            },
306            last_heartbeat_ms: chrono::Utc::now().timestamp_millis(),
307        })
308    }
309
310    fn validate_metadata_tags(
311        tags: &HashMap<String, String>,
312        encoded_bytes: usize,
313    ) -> Result<(), DiscoveryError> {
314        if tags.len() > MAX_METADATA_TAGS {
315            return Err(DiscoveryError::Serialization(format!(
316                "metadata tag count {} exceeds limit {MAX_METADATA_TAGS}",
317                tags.len()
318            )));
319        }
320        for (key, value) in tags {
321            if key.is_empty() || key.len() > MAX_METADATA_TAG_KEY_BYTES {
322                return Err(DiscoveryError::Serialization(format!(
323                    "metadata tag key must contain 1..={MAX_METADATA_TAG_KEY_BYTES} bytes"
324                )));
325            }
326            if value.len() > MAX_METADATA_TAG_VALUE_BYTES {
327                return Err(DiscoveryError::Serialization(format!(
328                    "metadata tag value for {key:?} exceeds {MAX_METADATA_TAG_VALUE_BYTES} bytes"
329                )));
330            }
331        }
332        if encoded_bytes > MAX_METADATA_TAGS_ENCODED_BYTES {
333            return Err(DiscoveryError::Serialization(format!(
334                "encoded metadata tags contain {encoded_bytes} bytes; limit is {MAX_METADATA_TAGS_ENCODED_BYTES}"
335            )));
336        }
337        Ok(())
338    }
339
340    fn validate_process_incarnation(tags: &HashMap<String, String>) -> Result<(), DiscoveryError> {
341        let valid = tags
342            .get(PROCESS_INCARNATION_TAG)
343            .and_then(|value| uuid::Uuid::parse_str(value).ok())
344            .is_some_and(|value| !value.is_nil());
345        if valid {
346            Ok(())
347        } else {
348            Err(DiscoveryError::Serialization(format!(
349                "metadata tag {PROCESS_INCARNATION_TAG:?} must contain a non-nil UUID"
350            )))
351        }
352    }
353
354    /// Build the chitchat key-value set for the local node.
355    fn local_kvs(info: &NodeInfo) -> Result<Vec<(String, String)>, DiscoveryError> {
356        if info.name.is_empty() || info.metadata.version != env!("CARGO_PKG_VERSION") {
357            return Err(DiscoveryError::Serialization(format!(
358                "gossip node name must be nonempty and version must be {:?}",
359                env!("CARGO_PKG_VERSION")
360            )));
361        }
362        Self::validate_metadata_tags(&info.metadata.tags, 0)?;
363        Self::validate_process_incarnation(&info.metadata.tags)?;
364        let canonical_tags: BTreeMap<&str, &str> = info
365            .metadata
366            .tags
367            .iter()
368            .map(|(key, value)| (key.as_str(), value.as_str()))
369            .collect();
370        let encoded_tags = serde_json::to_string(&canonical_tags)
371            .map_err(|error| DiscoveryError::Serialization(error.to_string()))?;
372        Self::validate_metadata_tags(&info.metadata.tags, encoded_tags.len())?;
373        let mut kvs = vec![
374            (keys::NODE_STATE.into(), info.state.to_string()),
375            (keys::RPC_ADDRESS.into(), info.rpc_address.clone()),
376            (keys::NODE_NAME.into(), info.name.clone()),
377            (keys::LOAD_CORES.into(), info.metadata.cores.to_string()),
378            (
379                keys::LOAD_MEMORY.into(),
380                info.metadata.memory_bytes.to_string(),
381            ),
382            (keys::NODE_VERSION.into(), info.metadata.version.clone()),
383            (
384                keys::PROTOCOL_VERSION.into(),
385                DISCOVERY_PROTOCOL_VERSION.into(),
386            ),
387            (keys::METADATA_TAGS.into(), encoded_tags),
388        ];
389        if let Some(ref fd) = info.metadata.failure_domain {
390            kvs.push((keys::FAILURE_DOMAIN.into(), fd.clone()));
391        }
392        Ok(kvs)
393    }
394}
395
396impl std::fmt::Debug for GossipDiscovery {
397    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
398        f.debug_struct("GossipDiscovery")
399            .field("config", &self.config)
400            .field("started", &self.started)
401            .finish_non_exhaustive()
402    }
403}
404
405impl GossipDiscovery {
406    /// Start with a caller-provided chitchat transport. Test harnesses
407    /// use this to inject a filtering / fault-injecting transport
408    /// wrapper (see
409    /// [`cluster::testing::PartitionableTransport`](crate::cluster::testing::PartitionableTransport)).
410    /// The regular [`Discovery::start`] just delegates here with a
411    /// default [`UdpTransport`](chitchat::transport::UdpTransport).
412    ///
413    /// # Errors
414    /// Same as [`Discovery::start`].
415    ///
416    /// # Panics
417    /// Panics via `unwrap` on an internal assertion if called twice
418    /// concurrently from the same `GossipDiscovery` — the `started`
419    /// flag check makes the second call a no-op.
420    #[allow(clippy::too_many_lines)]
421    pub async fn start_with_transport<T>(&mut self, transport: &T) -> Result<(), DiscoveryError>
422    where
423        T: chitchat::transport::Transport,
424    {
425        if self.started {
426            return Ok(());
427        }
428        let generation = self.config.process_generation;
429        if generation == 0 {
430            return Err(DiscoveryError::Serialization(
431                "gossip process generation must be nonzero".into(),
432            ));
433        }
434
435        let node_id = format!("node-{}", self.config.node_id.0);
436        let gossip_addr: std::net::SocketAddr = self
437            .config
438            .gossip_address
439            .parse()
440            .map_err(|e: std::net::AddrParseError| DiscoveryError::Bind(e.to_string()))?;
441
442        let advertise_addr = if let Some(ref host) = self.config.advertise_host {
443            let mut resolved = None;
444            #[cfg(feature = "cluster")]
445            {
446                if let Ok(addrs) = (host.as_str(), gossip_addr.port()).to_socket_addrs() {
447                    for addr in addrs {
448                        if addr.ip().is_ipv4() {
449                            resolved = Some(addr);
450                            break;
451                        }
452                    }
453                }
454            }
455            if let Some(addr) = resolved {
456                addr
457            } else {
458                return Err(DiscoveryError::Bind(format!(
459                    "failed to resolve configured advertise_host '{host}' (or cluster feature is disabled)"
460                )));
461            }
462        } else if gossip_addr.ip().is_unspecified() {
463            let resolved = {
464                let mut res = None;
465                #[cfg(feature = "cluster")]
466                {
467                    let hostname = gethostname::gethostname();
468                    let hostname_str = hostname.to_string_lossy();
469                    if !hostname_str.is_empty() {
470                        if let Ok(addrs) =
471                            (hostname_str.as_ref(), gossip_addr.port()).to_socket_addrs()
472                        {
473                            for addr in addrs {
474                                if addr.ip().is_ipv4() && !addr.ip().is_loopback() {
475                                    res = Some(addr);
476                                    break;
477                                }
478                            }
479                        }
480                    }
481                }
482                res
483            };
484            resolved.unwrap_or_else(|| {
485                std::net::SocketAddr::new(
486                    std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
487                    gossip_addr.port(),
488                )
489            })
490        } else {
491            gossip_addr
492        };
493
494        let seed_addrs: Vec<String> = self.config.seed_nodes.clone();
495
496        tracing::info!(
497            "Starting gossip discovery: gossip_addr = {}, advertise_addr = {}, seeds = {:?}",
498            gossip_addr,
499            advertise_addr,
500            seed_addrs
501        );
502
503        let config = chitchat::ChitchatConfig {
504            chitchat_id: chitchat::ChitchatId::new(node_id, generation, advertise_addr),
505            cluster_id: self.config.cluster_id.clone(),
506            gossip_interval: self.config.gossip_interval,
507            listen_addr: gossip_addr,
508            seed_nodes: seed_addrs,
509            failure_detector_config: chitchat::FailureDetectorConfig {
510                phi_threshold: self.config.phi_threshold,
511                initial_interval: self.config.gossip_interval,
512                // Map dead_node_grace_period to the failure detector's GC
513                // timer (W6 fix). Default is 24h which is far too long.
514                dead_node_grace_period: self.config.dead_node_grace_period,
515                ..Default::default()
516            },
517            marked_for_deletion_grace_period: self.config.dead_node_grace_period,
518            extra_liveness_predicate: None,
519            catchup_callback: None,
520        };
521
522        let initial_kvs = Self::local_kvs(&self.config.local_node)?;
523        let chitchat_handle = chitchat::spawn_chitchat(config, initial_kvs, transport)
524            .await
525            .map_err(|e| DiscoveryError::Bind(e.to_string()))?;
526
527        self.chitchat_handle = Some(chitchat_handle);
528
529        // Spawn membership watcher
530        let peers = Arc::clone(&self.peers);
531        let membership_tx = self.membership_tx.clone();
532        let cancel = self.cancel.clone();
533        let chitchat = self.chitchat_handle.as_ref().unwrap().chitchat().clone();
534        let local_node_id = self.config.node_id;
535
536        self.membership_handle = Some(tokio::spawn(async move {
537            let mut interval = tokio::time::interval(Duration::from_millis(500));
538            loop {
539                tokio::select! {
540                    biased;
541                    () = cancel.cancelled() => break,
542                    _ = interval.tick() => {
543                        let chitchat_guard = tokio::select! {
544                            biased;
545                            () = cancel.cancelled() => break,
546                            guard = chitchat.lock() => guard,
547                        };
548                        // Collect the set of live node IDs from the failure
549                        // detector so we only include reachable peers (C3 fix).
550                        let live_ids: std::collections::HashSet<&chitchat::ChitchatId> =
551                            chitchat_guard.live_nodes().collect();
552
553                        let nodes: Vec<_> = chitchat_guard.node_states().keys().map(|id| format!("{}(live={})", id.node_id, live_ids.contains(id))).collect();
554                        tracing::debug!("Chitchat state nodes: {:?}", nodes);
555
556                        // Select the unique highest durable process generation before parsing or
557                        // cloning any peer metadata. A malformed newest process must exclude its
558                        // stable node, not resurrect an older retained incarnation.
559                        let mut highest_generations = HashMap::new();
560                        for cc_id in chitchat_guard.node_states().keys() {
561                            let Some(node_id) = stable_node_id(&cc_id.node_id) else {
562                                continue;
563                            };
564                            if NodeId(node_id) != local_node_id {
565                                observe_peer_generation(
566                                    &mut highest_generations,
567                                    node_id,
568                                    cc_id.generation_id,
569                                );
570                            }
571                        }
572
573                        let mut new_peers = HashMap::new();
574                        for (cc_id, state) in chitchat_guard.node_states() {
575                            let Some(node_id) = stable_node_id(&cc_id.node_id) else {
576                                continue;
577                            };
578                            let Some(highest) = highest_generations.get(&node_id) else {
579                                continue;
580                            };
581                            if highest.ambiguous || cc_id.generation_id != highest.generation {
582                                continue;
583                            }
584
585                            let Some(encoded_tags) = state.get(keys::METADATA_TAGS) else {
586                                tracing::warn!(
587                                    node_id,
588                                    generation = cc_id.generation_id,
589                                    "excluding highest gossip generation without process identity"
590                                );
591                                continue;
592                            };
593                            if encoded_tags.len() > MAX_METADATA_TAGS_ENCODED_BYTES {
594                                tracing::warn!(
595                                    node_id,
596                                    generation = cc_id.generation_id,
597                                    "excluding peer with oversized gossip metadata tags"
598                                );
599                                continue;
600                            }
601                            let kvs: HashMap<String, String> = state
602                                .key_values()
603                                .filter(|(key, _)| is_node_info_key(key))
604                                .map(|(k, v)| (k.to_string(), v.to_string()))
605                                .collect();
606
607                            if let Some(mut info) = Self::parse_node_info(&cc_id.node_id, &kvs) {
608                                // Override self-reported state with failure detector opinion.
609                                if !live_ids.contains(cc_id) {
610                                    info.state = NodeState::Suspected;
611                                }
612                                new_peers.insert(node_id, info);
613                            } else {
614                                tracing::warn!(
615                                    node_id,
616                                    generation = cc_id.generation_id,
617                                    "excluding malformed highest gossip generation"
618                                );
619                            }
620                        }
621
622                        let peer_list: Vec<NodeInfo> =
623                            new_peers.values().cloned().collect();
624                        *peers.write() = new_peers;
625                        super::publish_if_changed(&membership_tx, peer_list);
626                    }
627                }
628            }
629        }));
630
631        self.started = true;
632        Ok(())
633    }
634
635    async fn stop_with_timeout(&mut self, timeout: Duration) {
636        self.cancel.cancel();
637        self.started = false;
638
639        let membership_handle = self.membership_handle.take();
640        let chitchat_handle = self.chitchat_handle.take();
641        let stop_membership = async move {
642            if let Some(handle) = membership_handle {
643                if let Some(Err(error)) =
644                    join_task_bounded(handle, timeout, "gossip-membership").await
645                {
646                    tracing::debug!(%error, "Gossip membership task stopped unexpectedly");
647                }
648            }
649        };
650        let stop_chitchat = async move {
651            if let Some(handle) = chitchat_handle {
652                let handle = ChitchatShutdownGuard::new(handle);
653                if let Err(error) = handle.handle().initiate_shutdown() {
654                    tracing::debug!(%error, "Chitchat server was already stopped");
655                }
656                match tokio::time::timeout(timeout, handle.handle().termination_watcher()).await {
657                    Ok(Ok(())) => {}
658                    Ok(Err(error)) => {
659                        tracing::warn!(%error, "Chitchat server stopped with an error");
660                    }
661                    Err(_) => {
662                        tracing::warn!(?timeout, "Chitchat server did not stop in time");
663                        handle.handle().abort();
664                        let _ = tokio::time::timeout(
665                            timeout.min(Duration::from_secs(1)),
666                            handle.handle().termination_watcher(),
667                        )
668                        .await;
669                    }
670                }
671            }
672        };
673        tokio::join!(stop_membership, stop_chitchat);
674
675        self.cancel = CancellationToken::new();
676    }
677}
678
679impl Discovery for GossipDiscovery {
680    async fn start(&mut self) -> Result<(), DiscoveryError> {
681        self.start_with_transport(&chitchat::transport::UdpTransport)
682            .await
683    }
684
685    // INVARIANT: discovery operations stay poll-driven for timeout and select cancellation.
686    #[allow(unknown_lints, clippy::unused_async, clippy::unused_async_trait_impl)]
687    async fn peers(&self) -> Result<Vec<NodeInfo>, DiscoveryError> {
688        if !self.started {
689            return Err(DiscoveryError::NotStarted);
690        }
691        let peers = self.peers.read();
692        Ok(peers.values().cloned().collect())
693    }
694
695    async fn announce(&self, info: NodeInfo) -> Result<(), DiscoveryError> {
696        if !self.started {
697            return Err(DiscoveryError::NotStarted);
698        }
699        if let Some(ref handle) = self.chitchat_handle {
700            let kvs = Self::local_kvs(&info)?;
701            handle
702                .with_chitchat(|chitchat| {
703                    for (key, value) in &kvs {
704                        chitchat.self_node_state().set(key.clone(), value.clone());
705                    }
706                })
707                .await;
708        }
709        Ok(())
710    }
711
712    fn membership_watch(&self) -> watch::Receiver<Vec<NodeInfo>> {
713        self.membership_rx.clone()
714    }
715
716    async fn stop(&mut self) -> Result<(), DiscoveryError> {
717        self.stop_with_timeout(DISCOVERY_SHUTDOWN_TIMEOUT).await;
718        Ok(())
719    }
720}
721
722impl Drop for GossipDiscovery {
723    fn drop(&mut self) {
724        self.cancel.cancel();
725        self.started = false;
726        if let Some(handle) = self.membership_handle.take() {
727            handle.abort();
728        }
729        if let Some(handle) = self.chitchat_handle.take() {
730            let _ = handle.initiate_shutdown();
731            handle.abort();
732        }
733    }
734}
735
736#[cfg(test)]
737mod tests;
laminar_core/cluster/discovery/gossip_discovery/mod.rs

laminar_core/cluster/discovery/gossip_discovery/
mod.rs
