1use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5
6use xds_core::NodeHash;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub struct StreamId(u64);
11
12impl StreamId {
13 pub fn new() -> Self {
15 static COUNTER: AtomicU64 = AtomicU64::new(1);
16 Self(COUNTER.fetch_add(1, Ordering::Relaxed))
17 }
18
19 #[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#[derive(Debug)]
46pub struct StreamContext {
47 id: StreamId,
49 node_hash: Option<NodeHash>,
51 node_id: Option<String>,
53 created_at: Instant,
55 requests: AtomicU64,
57 responses: AtomicU64,
59 last_request: std::sync::Mutex<Instant>,
61}
62
63impl StreamContext {
64 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 #[inline]
80 pub fn id(&self) -> StreamId {
81 self.id
82 }
83
84 #[inline]
86 pub fn node_hash(&self) -> Option<NodeHash> {
87 self.node_hash
88 }
89
90 #[inline]
92 pub fn node_id(&self) -> Option<&str> {
93 self.node_id.as_deref()
94 }
95
96 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 #[inline]
104 pub fn created_at(&self) -> Instant {
105 self.created_at
106 }
107
108 #[inline]
110 pub fn duration(&self) -> std::time::Duration {
111 self.created_at.elapsed()
112 }
113
114 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 pub fn record_response(&self) {
124 self.responses.fetch_add(1, Ordering::Relaxed);
125 }
126
127 #[inline]
129 pub fn request_count(&self) -> u64 {
130 self.requests.load(Ordering::Relaxed)
131 }
132
133 #[inline]
135 pub fn response_count(&self) -> u64 {
136 self.responses.load(Ordering::Relaxed)
137 }
138
139 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}