velo 0.3.0

Velo distributed-systems runtime: active messaging, peer discovery, streaming, rendezvous, and queue backends
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Peer registry for tracking remote peer state and handler availability.

use dashmap::DashMap;
use std::{collections::HashSet, sync::Arc, time::Instant};
use velo_ext::InstanceId;

/// State information for a registered peer.
#[derive(Clone, Debug)]
pub(crate) struct PeerState {
    /// Known handlers on this peer (None = haven't queried yet)
    pub handlers: Option<HashSet<String>>,
    /// Last time we communicated with this peer
    pub last_seen: Instant,
}

impl PeerState {
    pub fn new() -> Self {
        Self {
            handlers: None,
            last_seen: Instant::now(),
        }
    }

    pub fn update_last_seen(&mut self) {
        self.last_seen = Instant::now();
    }

    pub fn set_handlers(&mut self, handlers: Vec<String>) {
        self.handlers = Some(handlers.into_iter().collect());
        self.update_last_seen();
    }

    pub fn has_handler(&self, handler: &str) -> bool {
        self.handlers
            .as_ref()
            .map(|h| h.contains(handler))
            .unwrap_or(false)
    }
}

/// Registry for tracking remote peers and their state.
#[derive(Clone)]
pub(crate) struct PeerRegistry {
    peers: Arc<DashMap<InstanceId, PeerState>>,
}

impl PeerRegistry {
    pub fn new() -> Self {
        Self {
            peers: Arc::new(DashMap::new()),
        }
    }

    /// Register a peer without handler information.
    /// If the peer is already registered, preserves existing state (including cached handlers).
    pub fn register_peer(&self, instance_id: InstanceId) {
        self.peers.entry(instance_id).or_insert_with(PeerState::new);
    }

    /// Check if we have handler information for a peer.
    pub fn has_handler_info(&self, instance_id: InstanceId) -> bool {
        self.peers
            .get(&instance_id)
            .map(|state| state.handlers.is_some())
            .unwrap_or(false)
    }

    /// Get the list of handlers for a peer.
    pub fn get_handlers(&self, instance_id: InstanceId) -> Option<Vec<String>> {
        self.peers
            .get(&instance_id)
            .and_then(|state| state.handlers.as_ref().map(|h| h.iter().cloned().collect()))
    }

    /// Update the handler list for a peer.
    pub fn update_handlers(&self, instance_id: InstanceId, handlers: Vec<String>) {
        self.peers
            .entry(instance_id)
            .or_insert_with(PeerState::new)
            .set_handlers(handlers);
    }

    /// Check if a specific handler exists on a peer.
    pub fn handler_exists(&self, instance_id: InstanceId, handler: &str) -> bool {
        self.peers
            .get(&instance_id)
            .map(|state| state.has_handler(handler))
            .unwrap_or(false)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_peer_registry_basic() {
        let registry = PeerRegistry::new();
        let instance_id = InstanceId::new_v4();

        registry.register_peer(instance_id);
        assert!(!registry.has_handler_info(instance_id));

        registry.update_handlers(
            instance_id,
            vec!["handler1".to_string(), "handler2".to_string()],
        );
        assert!(registry.has_handler_info(instance_id));
        assert!(registry.handler_exists(instance_id, "handler1"));
        assert!(registry.handler_exists(instance_id, "handler2"));
        assert!(!registry.handler_exists(instance_id, "handler3"));

        let handlers = registry.get_handlers(instance_id).unwrap();
        assert_eq!(handlers.len(), 2);
        assert!(handlers.contains(&"handler1".to_string()));
        assert!(handlers.contains(&"handler2".to_string()));
    }

    #[test]
    fn test_peer_state_new() {
        let state = PeerState::new();
        assert!(state.handlers.is_none());
    }

    #[test]
    fn test_peer_state_set_handlers() {
        let mut state = PeerState::new();
        state.set_handlers(vec!["test".to_string()]);
        assert!(state.handlers.is_some());
        assert!(state.has_handler("test"));
        assert!(!state.has_handler("missing"));
    }
}