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