running_process/broker/
http_endpoint_registry.rs1use std::collections::HashMap;
15use std::sync::Mutex;
16
17use prost::Message;
18
19use crate::broker::protocol_v2::BackendHttpReady;
20
21pub type BackendId = String;
25
26#[derive(Debug, Default)]
32pub struct HttpEndpointRegistry {
33 inner: Mutex<HashMap<BackendId, Option<u16>>>,
34}
35
36impl HttpEndpointRegistry {
37 pub fn new() -> Self {
39 Self::default()
40 }
41
42 pub fn track(&self, backend_id: BackendId) {
44 let mut map = self.inner.lock().expect("registry mutex poisoned");
45 map.entry(backend_id).or_insert(None);
46 }
47
48 pub fn register_backend_http_endpoint(&self, backend_id: BackendId, port: u16) -> Option<u16> {
53 let mut map = self.inner.lock().expect("registry mutex poisoned");
54 map.insert(backend_id, Some(port)).flatten()
55 }
56
57 pub fn lookup(&self, backend_id: &str) -> Option<u16> {
63 let map = self.inner.lock().expect("registry mutex poisoned");
64 map.get(backend_id).copied().flatten()
65 }
66
67 pub fn state(&self, backend_id: &str) -> Option<Option<u16>> {
72 let map = self.inner.lock().expect("registry mutex poisoned");
73 map.get(backend_id).copied()
74 }
75
76 pub fn snapshot(&self) -> Vec<(BackendId, Option<u16>)> {
79 let map = self.inner.lock().expect("registry mutex poisoned");
80 map.iter().map(|(k, v)| (k.clone(), *v)).collect()
81 }
82}
83
84#[derive(Debug, thiserror::Error)]
86pub enum BackendHttpReadyError {
87 #[error("decode BackendHttpReady: {0}")]
89 Decode(#[from] prost::DecodeError),
90
91 #[error("BackendHttpReady.port = {0} is out of u16 range")]
93 PortOutOfRange(u32),
94}
95
96pub fn decode_and_register(
103 registry: &HttpEndpointRegistry,
104 backend_id: BackendId,
105 frame_body: &[u8],
106) -> Result<u16, BackendHttpReadyError> {
107 let ready = BackendHttpReady::decode(frame_body)?;
108 let port: u16 = ready
109 .port
110 .try_into()
111 .map_err(|_| BackendHttpReadyError::PortOutOfRange(ready.port))?;
112 registry.register_backend_http_endpoint(backend_id, port);
113 Ok(port)
114}
115
116#[cfg(test)]
117mod tests {
118 use super::*;
119
120 #[test]
121 fn empty_registry_has_no_state_for_unknown_backend() {
122 let reg = HttpEndpointRegistry::new();
123 assert!(reg.state("zccache").is_none());
124 assert!(reg.lookup("zccache").is_none());
125 }
126
127 #[test]
128 fn track_then_lookup_returns_none_for_pending_port() {
129 let reg = HttpEndpointRegistry::new();
130 reg.track("zccache".to_string());
131 assert!(reg.lookup("zccache").is_none());
133 assert_eq!(reg.state("zccache"), Some(None));
135 }
136
137 #[test]
138 fn register_endpoint_makes_port_available() {
139 let reg = HttpEndpointRegistry::new();
140 reg.track("zccache".to_string());
141 let prev = reg.register_backend_http_endpoint("zccache".to_string(), 8765);
142 assert_eq!(prev, None);
143 assert_eq!(reg.lookup("zccache"), Some(8765));
144 assert_eq!(reg.state("zccache"), Some(Some(8765)));
145 }
146
147 #[test]
148 fn register_endpoint_updates_existing_port_and_returns_previous() {
149 let reg = HttpEndpointRegistry::new();
150 reg.register_backend_http_endpoint("fbuild".to_string(), 8001);
151 let prev = reg.register_backend_http_endpoint("fbuild".to_string(), 8002);
152 assert_eq!(prev, Some(8001));
153 assert_eq!(reg.lookup("fbuild"), Some(8002));
154 }
155
156 #[test]
157 fn snapshot_reflects_all_tracked_backends() {
158 let reg = HttpEndpointRegistry::new();
159 reg.track("zccache".to_string());
160 reg.register_backend_http_endpoint("fbuild".to_string(), 8002);
161
162 let mut snap = reg.snapshot();
163 snap.sort();
164 assert_eq!(
165 snap,
166 vec![
167 ("fbuild".to_string(), Some(8002)),
168 ("zccache".to_string(), None)
169 ]
170 );
171 }
172
173 #[test]
174 fn decode_and_register_happy_path() {
175 let reg = HttpEndpointRegistry::new();
176 let msg = BackendHttpReady { port: 49_152 };
177 let mut body = Vec::with_capacity(msg.encoded_len());
178 msg.encode(&mut body).expect("encode BackendHttpReady");
179
180 let port = decode_and_register(®, "zccache".to_string(), &body)
181 .expect("decode_and_register succeeds");
182 assert_eq!(port, 49_152);
183 assert_eq!(reg.lookup("zccache"), Some(49_152));
184 }
185
186 #[test]
187 fn decode_and_register_rejects_oversized_port() {
188 let reg = HttpEndpointRegistry::new();
189 let msg = BackendHttpReady { port: 70_000 };
191 let mut body = Vec::with_capacity(msg.encoded_len());
192 msg.encode(&mut body).expect("encode BackendHttpReady");
193
194 let err = decode_and_register(®, "zccache".to_string(), &body)
195 .expect_err("port=70000 should be rejected");
196 match err {
197 BackendHttpReadyError::PortOutOfRange(70_000) => {}
198 other => panic!("expected PortOutOfRange(70000), got: {other:?}"),
199 }
200 assert!(reg.lookup("zccache").is_none());
202 }
203
204 #[test]
205 fn decode_and_register_rejects_malformed_frame() {
206 let reg = HttpEndpointRegistry::new();
207 let err = decode_and_register(®, "zccache".to_string(), &[0xFF; 8])
209 .expect_err("malformed frame should be rejected");
210 match err {
211 BackendHttpReadyError::Decode(_) => {}
212 other => panic!("expected Decode, got: {other:?}"),
213 }
214 assert!(reg.lookup("zccache").is_none());
215 }
216}