Skip to main content

rust_web_server/service_discovery/
mod.rs

1//! Dynamic backend pool with pluggable discovery sources.
2//!
3//! [`BackendPool`] maintains a thread-safe list of `"host:port"` addresses that
4//! can be refreshed on a background thread.  Discovery is delegated to a
5//! [`DiscoverySource`]:
6//!
7//! | Variant      | Description                                               |
8//! |------------- |---------------------------------------------------------- |
9//! | `Static`     | Fixed list — no polling required.                         |
10//! | `EnvPrefix`  | Scan `PREFIX_0`, `PREFIX_1`, … environment variables.     |
11//! | `File`       | Read one `host:port` per line from a file.                |
12//! | `Dns`        | A-record lookup — resolve hostname to all IPs.            |
13//! | `DnsSrv`     | SRV record lookup — weight-expanded `host:port` list.     |
14//! | `Consul`     | Consul HTTP API `/v1/health/service/:name`.                |
15//! | `Docker`     | Docker Engine API — containers carrying a given label.    |
16//! | `EtcdWatch`  | etcd v3 watch stream — incremental, low-latency updates.   |
17//!
18//! # Example
19//!
20//! ```rust,no_run
21//! use rust_web_server::service_discovery::BackendPool;
22//!
23//! // Fixed list — no background thread needed.
24//! let pool = BackendPool::r#static(vec!["10.0.0.1:8080".into(), "10.0.0.2:8080".into()]);
25//! println!("{:?}", pool.backends());
26//!
27//! // Env-var discovery, refreshed every 60 seconds.
28//! let pool = BackendPool::env_prefix("MY_SVC_BACKEND")
29//!     .poll_interval_secs(60);
30//! pool.start();
31//! println!("{:?}", pool.backends());
32//! ```
33
34#[cfg(test)]
35mod tests;
36// `pub(crate)` (not just `mod`) so `secrets`'s Vault/AWS-SM/Key-Vault backends
37// can reuse this parser for their own small JSON responses, instead of a
38// third hand-rolled JSON parser alongside this one and `mcp::json_rpc`.
39pub(crate) mod json_lite;
40mod consul;
41mod dns_srv;
42mod docker;
43mod etcd;
44
45use std::net::ToSocketAddrs;
46use std::sync::{Arc, RwLock};
47use std::time::Duration;
48
49// ── DiscoverySource ───────────────────────────────────────────────────────────
50
51/// Controls how [`BackendPool`] discovers backend addresses.
52pub enum DiscoverySource {
53    /// Fixed list of `"host:port"` addresses — never refreshed.
54    Static(Vec<String>),
55    /// Scan environment variables `PREFIX_0`, `PREFIX_1`, … until one is absent.
56    EnvPrefix(String),
57    /// Read one `host:port` per line from a file.  Blank lines and lines starting
58    /// with `#` are ignored.
59    File(String),
60    /// Resolve `hostname` via A-record DNS lookup; format each IP as `ip:port`.
61    Dns { hostname: String, port: u16 },
62    /// Resolve `record` (e.g. `_http._tcp.example.com`) via SRV lookup. Only the
63    /// lowest-priority tier of records is used; within that tier each `target:port`
64    /// is repeated `weight.clamp(1, 20)` times in the returned list so a plain
65    /// round-robin consumer still sees proportional selection frequency — see
66    /// [`dns_srv`] module docs for the exact algorithm and its bound.
67    DnsSrv { record: String },
68    /// Query a Consul agent's `/v1/health/service/:name` endpoint. `addr` is
69    /// `host:port` of the agent (e.g. `127.0.0.1:8500`); only instances passing
70    /// all health checks are returned.
71    Consul { addr: String, service: String },
72    /// Query the Docker Engine API (over a Unix domain socket) for running
73    /// containers carrying `label`, using each container's *value* for that
74    /// label as the backend address directly (e.g. `rws.backend=10.0.0.5:8080`)
75    /// — see [`docker`] module docs for why the label value, not a published
76    /// port guess, is the address. Unix-only; a no-op elsewhere.
77    Docker { label: String, socket_path: String },
78    /// Watch an etcd v3 key prefix via its gRPC-gateway JSON/HTTP `/v3/watch`
79    /// endpoint. Unlike every other source, this is *not* driven by the
80    /// generic poll loop — [`BackendPool::start`] spawns a dedicated
81    /// long-lived connection instead, applying `PUT`/`DELETE` events
82    /// incrementally as they arrive rather than re-listing on each one.
83    /// `resolve()` (and therefore plain `refresh()`) still performs a one-shot
84    /// `/v3/kv/range` listing, so a pool built with this source is usable
85    /// without ever calling `start()`, just without live updates. Plain HTTP
86    /// only — no TLS support yet.
87    EtcdWatch { endpoints: Vec<String>, prefix: String },
88}
89
90impl DiscoverySource {
91    /// Perform a single discovery cycle and return the current backend list.
92    fn resolve(&self) -> Vec<String> {
93        match self {
94            DiscoverySource::Static(v) => v.clone(),
95
96            DiscoverySource::EnvPrefix(prefix) => {
97                let mut backends = Vec::new();
98                let mut i = 0usize;
99                loop {
100                    let key = format!("{}_{}", prefix, i);
101                    match std::env::var(&key) {
102                        Ok(val) => { backends.push(val); i += 1; }
103                        Err(_) => break,
104                    }
105                }
106                backends
107            }
108
109            DiscoverySource::File(path) => {
110                match std::fs::read_to_string(path) {
111                    Ok(contents) => contents
112                        .lines()
113                        .map(str::trim)
114                        .filter(|line| !line.is_empty() && !line.starts_with('#'))
115                        .map(str::to_string)
116                        .collect(),
117                    Err(e) => {
118                        eprintln!("service_discovery: cannot read backend file {:?}: {}", path, e);
119                        Vec::new()
120                    }
121                }
122            }
123
124            DiscoverySource::Dns { hostname, port } => {
125                let addr_str = format!("{}:{}", hostname, port);
126                match addr_str.to_socket_addrs() {
127                    Ok(addrs) => addrs
128                        .map(|sa| format!("{}:{}", sa.ip(), sa.port()))
129                        .collect(),
130                    Err(e) => {
131                        eprintln!("service_discovery: DNS lookup for {} failed: {}", addr_str, e);
132                        Vec::new()
133                    }
134                }
135            }
136
137            DiscoverySource::DnsSrv { record } => dns_srv::resolve(record),
138
139            DiscoverySource::Consul { addr, service } => consul::discover(addr, service),
140
141            DiscoverySource::Docker { label, socket_path } => docker::discover(socket_path, label),
142
143            DiscoverySource::EtcdWatch { endpoints, prefix } => {
144                let Some(endpoint) = endpoints.first() else { return Vec::new() };
145                match etcd::kv_range(endpoint, prefix) {
146                    Ok(map) => map.into_values().collect(),
147                    Err(e) => {
148                        eprintln!("service_discovery: etcd range query failed: {}", e);
149                        Vec::new()
150                    }
151                }
152            }
153        }
154    }
155}
156
157// ── BackendPool ───────────────────────────────────────────────────────────────
158
159/// Thread-safe pool of backend addresses, optionally refreshed in the background.
160///
161/// Clone this type freely — all clones share the same underlying `RwLock<Vec>`.
162#[derive(Clone)]
163pub struct BackendPool {
164    backends: Arc<RwLock<Vec<String>>>,
165    source: Arc<DiscoverySource>,
166    poll_interval_secs: u64,
167}
168
169impl BackendPool {
170    fn new(source: DiscoverySource) -> Self {
171        Self {
172            backends: Arc::new(RwLock::new(Vec::new())),
173            source: Arc::new(source),
174            poll_interval_secs: 30,
175        }
176    }
177
178    /// Create a pool from a fixed list of backends.
179    ///
180    /// The list is available immediately; `start()` is a no-op.
181    pub fn r#static(backends: Vec<String>) -> Self {
182        let initial = backends.clone();
183        let pool = Self::new(DiscoverySource::Static(backends));
184        *pool.backends.write().unwrap() = initial;
185        pool
186    }
187
188    /// Create a pool whose backends are read from environment variables
189    /// `PREFIX_0`, `PREFIX_1`, … at startup and every `poll_interval_secs`.
190    pub fn env_prefix(prefix: impl Into<String>) -> Self {
191        Self::new(DiscoverySource::EnvPrefix(prefix.into()))
192    }
193
194    /// Create a pool whose backends are read from a file (one `host:port` per line).
195    pub fn file(path: impl Into<String>) -> Self {
196        Self::new(DiscoverySource::File(path.into()))
197    }
198
199    /// Create a pool whose backends are discovered via DNS A-record lookup.
200    pub fn dns(hostname: impl Into<String>, port: u16) -> Self {
201        Self::new(DiscoverySource::Dns { hostname: hostname.into(), port })
202    }
203
204    /// Create a pool whose backends are discovered via DNS SRV lookup (e.g.
205    /// `_http._tcp.my-service.default.svc.cluster.local` for a Kubernetes
206    /// headless Service). See [`DiscoverySource::DnsSrv`] for the
207    /// priority/weight handling.
208    pub fn dns_srv(record: impl Into<String>) -> Self {
209        Self::new(DiscoverySource::DnsSrv { record: record.into() })
210    }
211
212    /// Create a pool whose backends are discovered via a Consul agent's
213    /// `/v1/health/service/:name` endpoint. `addr` is `host:port` of the
214    /// agent, e.g. `"127.0.0.1:8500"`. Only instances passing all health
215    /// checks are returned.
216    pub fn consul(addr: impl Into<String>, service: impl Into<String>) -> Self {
217        Self::new(DiscoverySource::Consul { addr: addr.into(), service: service.into() })
218    }
219
220    /// Create a pool whose backends are discovered from running Docker
221    /// containers carrying `label`, using the default socket path
222    /// (`/var/run/docker.sock`). See [`DiscoverySource::Docker`] for the
223    /// label-value-is-the-address convention. Unix-only; a no-op elsewhere.
224    pub fn docker(label: impl Into<String>) -> Self {
225        Self::new(DiscoverySource::Docker {
226            label: label.into(),
227            socket_path: "/var/run/docker.sock".to_string(),
228        })
229    }
230
231    /// Same as [`BackendPool::docker`] but against a non-default Docker
232    /// socket path.
233    pub fn docker_with_socket(label: impl Into<String>, socket_path: impl Into<String>) -> Self {
234        Self::new(DiscoverySource::Docker { label: label.into(), socket_path: socket_path.into() })
235    }
236
237    /// Create a pool whose backends are discovered from an etcd v3 key
238    /// `prefix`, kept live via a watch stream once [`BackendPool::start`] is
239    /// called. `endpoints` is a list of `host:port` etcd cluster members —
240    /// only the first is used today (no client-side failover yet). Plain
241    /// HTTP only. See [`DiscoverySource::EtcdWatch`].
242    pub fn etcd(endpoints: Vec<String>, prefix: impl Into<String>) -> Self {
243        Self::new(DiscoverySource::EtcdWatch { endpoints, prefix: prefix.into() })
244    }
245
246    /// Override the background refresh interval (default: 30 seconds).
247    ///
248    /// Only meaningful for `File` and `Dns` sources.
249    pub fn poll_interval_secs(mut self, secs: u64) -> Self {
250        self.poll_interval_secs = secs;
251        self
252    }
253
254    /// Start the background refresh thread.
255    ///
256    /// For `Static` sources this is a no-op. For all others, an immediate
257    /// `refresh()` is performed before spawning the background thread so that
258    /// the first `backends()` call returns a populated list.
259    ///
260    /// `EtcdWatch` is special-cased: instead of the generic sleep-and-poll
261    /// loop, a dedicated thread holds a long-lived connection to etcd's watch
262    /// stream and applies `PUT`/`DELETE` events to the backend list as they
263    /// arrive — no polling interval involved after the initial `refresh()`.
264    pub fn start(&self) {
265        if matches!(self.source.as_ref(), DiscoverySource::Static(_)) {
266            return;
267        }
268        self.refresh();
269
270        if let DiscoverySource::EtcdWatch { endpoints, prefix } = self.source.as_ref() {
271            let endpoints = endpoints.clone();
272            let prefix = prefix.clone();
273            let pool = self.clone();
274            std::thread::spawn(move || etcd::watch_forever(&endpoints, &prefix, &pool));
275            return;
276        }
277
278        let pool = self.clone();
279        let interval = Duration::from_secs(self.poll_interval_secs);
280        std::thread::spawn(move || loop {
281            std::thread::sleep(interval);
282            pool.refresh();
283        });
284    }
285
286    /// Return a snapshot of the current backend list.
287    pub fn backends(&self) -> Vec<String> {
288        self.backends.read().unwrap().clone()
289    }
290
291    /// Replace the current backend list with `backends`.
292    ///
293    /// Useful for testing or for external control planes that push updates.
294    pub fn update(&self, backends: Vec<String>) {
295        *self.backends.write().unwrap() = backends;
296    }
297
298    /// Perform one synchronous refresh cycle.
299    ///
300    /// Called automatically by `start()` and by the background thread.
301    pub fn refresh(&self) {
302        let resolved = self.source.resolve();
303        *self.backends.write().unwrap() = resolved;
304    }
305}