Skip to main content

xds_server/
stream.rs

1//! Stream context and identification.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5
6use xds_core::NodeHash;
7
8/// Unique identifier for a stream.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub struct StreamId(u64);
11
12impl StreamId {
13    /// Generate a new unique stream ID.
14    pub fn new() -> Self {
15        static COUNTER: AtomicU64 = AtomicU64::new(1);
16        Self(COUNTER.fetch_add(1, Ordering::Relaxed))
17    }
18
19    /// Get the numeric value.
20    #[inline]
21    pub fn as_u64(&self) -> u64 {
22        self.0
23    }
24}
25
26impl Default for StreamId {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl std::fmt::Display for StreamId {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        write!(f, "stream-{}", self.0)
35    }
36}
37
38/// Context for an active xDS stream.
39///
40/// Tracks metadata about a client connection including:
41/// - Stream identifier
42/// - Node information
43/// - Timing information
44/// - Request/response counts
45#[derive(Debug)]
46pub struct StreamContext {
47    /// Unique stream identifier.
48    id: StreamId,
49    /// Node hash for this stream's client.
50    node_hash: Option<NodeHash>,
51    /// Node ID as a string.
52    node_id: Option<String>,
53    /// When the stream was created.
54    created_at: Instant,
55    /// Number of requests received.
56    requests: AtomicU64,
57    /// Number of responses sent.
58    responses: AtomicU64,
59    /// Last request timestamp.
60    last_request: std::sync::Mutex<Instant>,
61}
62
63impl StreamContext {
64    /// Create a new stream context.
65    pub fn new() -> Self {
66        let now = Instant::now();
67        Self {
68            id: StreamId::new(),
69            node_hash: None,
70            node_id: None,
71            created_at: now,
72            requests: AtomicU64::new(0),
73            responses: AtomicU64::new(0),
74            last_request: std::sync::Mutex::new(now),
75        }
76    }
77
78    /// Get the stream ID.
79    #[inline]
80    pub fn id(&self) -> StreamId {
81        self.id
82    }
83
84    /// Get the node hash if set.
85    #[inline]
86    pub fn node_hash(&self) -> Option<NodeHash> {
87        self.node_hash
88    }
89
90    /// Get the node ID if set.
91    #[inline]
92    pub fn node_id(&self) -> Option<&str> {
93        self.node_id.as_deref()
94    }
95
96    /// Set the node information.
97    pub fn set_node(&mut self, node_id: String, node_hash: NodeHash) {
98        self.node_id = Some(node_id);
99        self.node_hash = Some(node_hash);
100    }
101
102    /// Get when this stream was created.
103    #[inline]
104    pub fn created_at(&self) -> Instant {
105        self.created_at
106    }
107
108    /// Get stream duration.
109    #[inline]
110    pub fn duration(&self) -> std::time::Duration {
111        self.created_at.elapsed()
112    }
113
114    /// Record a request.
115    pub fn record_request(&self) {
116        self.requests.fetch_add(1, Ordering::Relaxed);
117        if let Ok(mut last) = self.last_request.lock() {
118            *last = Instant::now();
119        }
120    }
121
122    /// Record a response.
123    pub fn record_response(&self) {
124        self.responses.fetch_add(1, Ordering::Relaxed);
125    }
126
127    /// Get total requests.
128    #[inline]
129    pub fn request_count(&self) -> u64 {
130        self.requests.load(Ordering::Relaxed)
131    }
132
133    /// Get total responses.
134    #[inline]
135    pub fn response_count(&self) -> u64 {
136        self.responses.load(Ordering::Relaxed)
137    }
138
139    /// Get time since last request.
140    pub fn idle_time(&self) -> std::time::Duration {
141        self.last_request
142            .lock()
143            .map(|t| t.elapsed())
144            .unwrap_or_default()
145    }
146}
147
148impl Default for StreamContext {
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn stream_id_unique() {
160        let id1 = StreamId::new();
161        let id2 = StreamId::new();
162        assert_ne!(id1, id2);
163    }
164
165    #[test]
166    fn stream_context_basic() {
167        let ctx = StreamContext::new();
168        assert_eq!(ctx.request_count(), 0);
169        assert_eq!(ctx.response_count(), 0);
170        assert!(ctx.node_hash().is_none());
171    }
172
173    #[test]
174    fn stream_context_counting() {
175        let ctx = StreamContext::new();
176        ctx.record_request();
177        ctx.record_request();
178        ctx.record_response();
179
180        assert_eq!(ctx.request_count(), 2);
181        assert_eq!(ctx.response_count(), 1);
182    }
183
184    #[test]
185    fn stream_context_node() {
186        let mut ctx = StreamContext::new();
187        let hash = NodeHash::from_id("test-node");
188        ctx.set_node("test-node".to_string(), hash);
189
190        assert_eq!(ctx.node_id(), Some("test-node"));
191        assert_eq!(ctx.node_hash(), Some(hash));
192    }
193}