Skip to main content

dynomite/cluster/
pool.rs

1//! Server pool owner.
2//!
3//! [`ServerPool`] is the cluster-level container that holds the
4//! local datastore endpoint, the dnode listener configuration, the
5//! datacenter and rack tables, the per-peer connection pools, the
6//! consistency settings, and the response-manager constructors used
7//! by [`crate::cluster::dispatch::ClusterDispatcher`]. It holds the
8//! server-pool data shape and is the entry point that the
9//! `dynomited` binary and the embedding API construct from a parsed
10//! [`crate::conf::Config`].
11
12use std::sync::Arc;
13use std::time::Duration;
14
15use parking_lot::RwLock;
16
17use crate::cluster::datacenter::Datacenter;
18use crate::cluster::peer::Peer;
19use crate::conf::{
20    ConfPool, ConsistencyLevel as ConfConsistencyLevel, DataStore, Distribution, HashType,
21};
22
23use crate::msg::ConsistencyLevel;
24use crate::net::auto_eject::AutoEject;
25
26/// Convert a YAML-form consistency string from a bucket-type
27/// stanza into the runtime [`ConsistencyLevel`] enum.
28///
29/// Values that fail validation fall back to `DcOne`; the
30/// configuration validator runs first, so by the time this is
31/// called every string is already known to be one of the four
32/// canonical names (case-insensitive).
33fn bucket_type_consistency(raw: &str) -> ConsistencyLevel {
34    match ConfConsistencyLevel::parse("bucket_type_consistency", raw) {
35        Ok(ConfConsistencyLevel::DcQuorum) => ConsistencyLevel::DcQuorum,
36        Ok(ConfConsistencyLevel::DcSafeQuorum) => ConsistencyLevel::DcSafeQuorum,
37        Ok(ConfConsistencyLevel::DcEachSafeQuorum) => ConsistencyLevel::DcEachSafeQuorum,
38        Ok(ConfConsistencyLevel::DcOne) | Err(_) => ConsistencyLevel::DcOne,
39    }
40}
41
42/// Resolved bucket-type bundle stored on the live
43/// [`PoolConfig`].
44///
45/// Mirrors [`crate::conf::ConfBucketType`] but with the
46/// consistency strings already parsed into
47/// [`ConsistencyLevel`] and the field-name semantics finalised
48/// for the dispatcher's hot path.
49#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct BucketType {
51    /// Bucket name. Compared verbatim against the prefix
52    /// returned by [`crate::proto::redis::bucket_name`].
53    pub name: String,
54    /// Read consistency override.
55    pub read_consistency: ConsistencyLevel,
56    /// Write consistency override.
57    pub write_consistency: ConsistencyLevel,
58    /// Replica fan-out cap. `0` means no cap.
59    pub n_val: u8,
60}
61
62/// Minimal projection of the YAML pool block consumed by the
63/// cluster runtime.
64///
65/// Carries the fields copied from
66/// `ConfPool` into the runtime pool during initialization.
67#[derive(Clone, Debug)]
68pub struct PoolConfig {
69    /// Pool name.
70    pub name: String,
71    /// Local datacenter name.
72    pub dc: String,
73    /// Local rack name.
74    pub rack: String,
75    /// Backing datastore protocol.
76    pub data_store: DataStore,
77    /// Hash function used for token ring lookups.
78    pub hash: HashType,
79    /// Distribution algorithm used to map a hashed key to a
80    /// peer. Defaults to [`Distribution::Vnode`].
81    pub distribution: Distribution,
82    /// Optional shadow distribution. When `Some`, the
83    /// dispatcher computes both the live and shadow routes for
84    /// every request and bumps a counter when they disagree.
85    /// The actual route is the configured
86    /// [`Self::distribution`].
87    pub distribution_shadow: Option<Distribution>,
88    /// Read consistency level.
89    pub read_consistency: ConsistencyLevel,
90    /// Write consistency level.
91    pub write_consistency: ConsistencyLevel,
92    /// Operation timeout in milliseconds.
93    pub timeout_ms: u64,
94    /// Eject window (`server_retry_timeout_ms`).
95    pub server_retry_timeout_ms: u64,
96    /// Consecutive-failure threshold.
97    pub server_failure_limit: u32,
98    /// Honor `auto_eject_hosts`.
99    pub auto_eject_hosts: bool,
100    /// Whether gossip is enabled (`enable_gossip`).
101    pub enable_gossip: bool,
102    /// Per-bucket routing-property bundles. Empty when the
103    /// `bucket_types:` YAML stanza is unset.
104    pub bucket_types: Vec<BucketType>,
105    /// Name of the bucket type to apply when the request key has
106    /// no slash. References an entry of `bucket_types`; `None`
107    /// keeps the pool-level defaults.
108    pub default_bucket_type: Option<String>,
109    /// Hinted-handoff feature flag. When `true`, writes targeted
110    /// at peers in [`crate::cluster::peer::PeerState::Down`] (or
111    /// at peers whose outbound channel is closed / full) are
112    /// stored as hints in the node-local
113    /// [`crate::cluster::hints::HintStore`] and counted toward
114    /// the request's consistency threshold. The default `false`
115    /// preserves the legacy behaviour where Down targets are
116    /// silently skipped.
117    pub enable_hinted_handoff: bool,
118    /// Hint expiry, in seconds. Hints older than this are
119    /// dropped during the periodic sweep so the in-memory store
120    /// stays bounded.
121    pub hint_ttl_seconds: u64,
122    /// Upper bound on the in-memory hint store. Once the store
123    /// holds this many bytes, further enqueues fail and the
124    /// dispatcher falls back to its non-handoff error path.
125    pub hint_store_max_bytes: u64,
126    /// Hint drainer cadence, in milliseconds. Setting this to
127    /// zero is rejected by the configuration validator when
128    /// `enable_hinted_handoff` is true.
129    pub hint_drain_interval_ms: u64,
130    /// Optional directory for the durable hinted-handoff backend.
131    /// When `Some` (and `enable_hinted_handoff` is true), the
132    /// hint store persists one segment file per peer here and
133    /// replays them at startup. `None` keeps the store RAM-only.
134    pub hint_dir: Option<std::path::PathBuf>,
135}
136
137impl Default for PoolConfig {
138    fn default() -> Self {
139        Self {
140            name: "p".into(),
141            dc: "localdc".into(),
142            rack: "localrack".into(),
143            data_store: DataStore::Valkey,
144            hash: HashType::Murmur,
145            distribution: Distribution::Vnode,
146            distribution_shadow: None,
147            read_consistency: ConsistencyLevel::DcOne,
148            write_consistency: ConsistencyLevel::DcOne,
149            timeout_ms: 5_000,
150            server_retry_timeout_ms: 30_000,
151            server_failure_limit: 2,
152            auto_eject_hosts: false,
153            enable_gossip: false,
154            bucket_types: Vec::new(),
155            default_bucket_type: None,
156            enable_hinted_handoff: false,
157            hint_ttl_seconds: 86_400,
158            hint_store_max_bytes: 64 * 1024 * 1024,
159            hint_drain_interval_ms: 30_000,
160            hint_dir: None,
161        }
162    }
163}
164
165impl PoolConfig {
166    /// Look up a bucket type by name.
167    ///
168    /// Returns the matching [`BucketType`] when one is
169    /// configured. The dispatcher consults this on every request
170    /// to swap in per-bucket consistency / fan-out settings.
171    #[must_use]
172    pub fn lookup_bucket_type(&self, name: &[u8]) -> Option<&BucketType> {
173        self.bucket_types
174            .iter()
175            .find(|bt| bt.name.as_bytes() == name)
176    }
177
178    /// Resolve the bucket type that applies to a request whose
179    /// extracted bucket name is `bucket`. `None` falls back to
180    /// `default_bucket_type` (also possibly `None`).
181    #[must_use]
182    pub fn resolve_bucket_type(&self, bucket: Option<&[u8]>) -> Option<&BucketType> {
183        if let Some(b) = bucket {
184            if let Some(bt) = self.lookup_bucket_type(b) {
185                return Some(bt);
186            }
187        }
188        // Either no bucket prefix on the key, or the prefix did
189        // not match any configured type. Fall through to the
190        // default bucket type when one is named.
191        let name = self.default_bucket_type.as_deref()?;
192        self.lookup_bucket_type(name.as_bytes())
193    }
194}
195
196impl PoolConfig {
197    /// Construct a `PoolConfig` from a [`ConfPool`] block. Fields
198    /// missing from the YAML are filled with defaults (the
199    /// caller is expected to have called
200    /// [`crate::conf::Config::finalize`] before this point).
201    ///
202    /// # Examples
203    ///
204    /// ```
205    /// use dynomite::cluster::pool::PoolConfig;
206    /// use dynomite::conf::Config;
207    /// let mut cfg = Config::parse_str(
208    ///     "p:\n  listen: 127.0.0.1:1\n  dyn_listen: 127.0.0.1:2\n  tokens: '1'\n  servers:\n  - 127.0.0.1:3:1\n  data_store: 0\n",
209    /// ).unwrap();
210    /// cfg.finalize();
211    /// let pc = PoolConfig::from_conf("p", cfg.pool());
212    /// assert_eq!(pc.name, "p");
213    /// ```
214    #[must_use]
215    pub fn from_conf(name: &str, pool: &ConfPool) -> Self {
216        let parse_consistency = |s: &Option<String>| {
217            s.as_deref()
218                .and_then(ConsistencyLevel::from_name)
219                .unwrap_or(ConsistencyLevel::DcOne)
220        };
221        let data_store = match pool.data_store {
222            Some(1) => DataStore::Memcache,
223            Some(2) => DataStore::Dyniak,
224            _ => DataStore::Valkey,
225        };
226        Self {
227            name: name.to_string(),
228            dc: pool.datacenter.clone().unwrap_or_else(|| "localdc".into()),
229            rack: pool.rack.clone().unwrap_or_else(|| "localrack".into()),
230            data_store,
231            hash: pool.hash.unwrap_or(HashType::Murmur),
232            distribution: pool.resolved_distribution(),
233            distribution_shadow: pool.distribution_shadow,
234            read_consistency: parse_consistency(&pool.read_consistency),
235            write_consistency: parse_consistency(&pool.write_consistency),
236            timeout_ms: pool
237                .timeout
238                .and_then(|n| u64::try_from(n).ok())
239                .unwrap_or(5_000),
240            server_retry_timeout_ms: pool
241                .server_retry_timeout
242                .and_then(|n| u64::try_from(n).ok())
243                .unwrap_or(30_000),
244            server_failure_limit: pool
245                .server_failure_limit
246                .and_then(|n| u32::try_from(n).ok())
247                .unwrap_or(2),
248            auto_eject_hosts: pool.auto_eject_hosts.unwrap_or(false),
249            enable_gossip: pool.enable_gossip.unwrap_or(false),
250            bucket_types: pool
251                .bucket_types
252                .iter()
253                .map(|bt| BucketType {
254                    name: bt.name.clone(),
255                    read_consistency: bucket_type_consistency(&bt.read_consistency),
256                    write_consistency: bucket_type_consistency(&bt.write_consistency),
257                    n_val: bt.n_val,
258                })
259                .collect(),
260            default_bucket_type: pool.default_bucket_type.clone(),
261            enable_hinted_handoff: pool.enable_hinted_handoff.unwrap_or(false),
262            hint_ttl_seconds: pool.hint_ttl_seconds.unwrap_or(86_400),
263            hint_store_max_bytes: pool.hint_store_max_bytes.unwrap_or(64 * 1024 * 1024),
264            hint_drain_interval_ms: pool.hint_drain_interval_ms.unwrap_or(30_000),
265            hint_dir: pool.hint_dir.clone(),
266        }
267    }
268}
269
270/// Cluster-wide owner.
271///
272/// Holds the topology (datacenters, racks), the peer list (peer
273/// index 0 is always the local node, mirroring the reference
274/// engine), and the per-peer auto-eject decision state.
275///
276/// `peers` and `datacenters` live behind `RwLock`s so the
277/// dispatcher can hold a read lock while gossip occasionally
278/// upgrades to write.
279///
280/// # Examples
281///
282/// ```
283/// use dynomite::cluster::pool::{PoolConfig, ServerPool};
284/// use dynomite::cluster::peer::{Peer, PeerEndpoint};
285/// use dynomite::hashkit::DynToken;
286/// let cfg = PoolConfig {
287///     dc: "dc1".into(),
288///     rack: "r1".into(),
289///     ..PoolConfig::default()
290/// };
291/// let local = Peer::new(
292///     0, PeerEndpoint::tcp("127.0.0.1".into(), 8101), "r1".into(), "dc1".into(),
293///     vec![DynToken::from_u32(1)], true, true, false,
294/// );
295/// let pool = ServerPool::new(cfg, vec![local]);
296/// assert_eq!(pool.peers().read().len(), 1);
297/// ```
298#[derive(Debug)]
299pub struct ServerPool {
300    config: PoolConfig,
301    peers: Arc<RwLock<Vec<Peer>>>,
302    datacenters: Arc<RwLock<Vec<Datacenter>>>,
303    auto_eject: Arc<RwLock<Vec<AutoEject>>>,
304}
305
306impl ServerPool {
307    /// Build a fresh pool from a [`PoolConfig`] and an initial peer
308    /// list (peer index 0 is the local node).
309    ///
310    /// Datacenters and racks are populated automatically from the
311    /// supplied peers; their continuum is rebuilt by
312    /// [`ServerPool::rebuild_ring`].
313    ///
314    /// # Examples
315    ///
316    /// ```
317    /// use dynomite::cluster::pool::{PoolConfig, ServerPool};
318    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
319    /// # use dynomite::hashkit::DynToken;
320    /// # use dynomite::conf::{DataStore, HashType};
321    /// # use dynomite::msg::ConsistencyLevel;
322    /// # let cfg = PoolConfig::default();
323    /// # let local = Peer::new(
324    /// #    0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
325    /// #    vec![DynToken::from_u32(0)], true, true, false,
326    /// # );
327    /// let pool = ServerPool::new(cfg, vec![local]);
328    /// pool.rebuild_ring();
329    /// assert_eq!(pool.datacenters().read().len(), 1);
330    /// ```
331    #[must_use]
332    pub fn new(config: PoolConfig, peers: Vec<Peer>) -> Self {
333        let mut dcs: Vec<Datacenter> = Vec::new();
334        for p in &peers {
335            let dc_idx = if let Some(i) = dcs.iter().position(|d| d.name() == p.dc()) {
336                i
337            } else {
338                dcs.push(Datacenter::new(p.dc().to_string()));
339                dcs.len() - 1
340            };
341            dcs[dc_idx].upsert_rack(p.rack().to_string());
342        }
343        let auto_eject_template = AutoEject::new(
344            config.auto_eject_hosts,
345            config.server_failure_limit,
346            Duration::from_millis(config.server_retry_timeout_ms),
347        );
348        let mut auto_ejects = Vec::with_capacity(peers.len());
349        for _ in &peers {
350            auto_ejects.push(auto_eject_template.clone());
351        }
352        let pool = Self {
353            config,
354            peers: Arc::new(RwLock::new(peers)),
355            datacenters: Arc::new(RwLock::new(dcs)),
356            auto_eject: Arc::new(RwLock::new(auto_ejects)),
357        };
358        pool.rebuild_ring();
359        pool
360    }
361
362    /// Configuration block.
363    #[must_use]
364    pub fn config(&self) -> &PoolConfig {
365        &self.config
366    }
367
368    /// Borrow the peer list (RwLock).
369    #[must_use]
370    pub fn peers(&self) -> &RwLock<Vec<Peer>> {
371        &self.peers
372    }
373
374    /// Shared `Arc` to the peer list.
375    #[must_use]
376    pub fn peers_arc(&self) -> Arc<RwLock<Vec<Peer>>> {
377        self.peers.clone()
378    }
379
380    /// Borrow the datacenter list.
381    #[must_use]
382    pub fn datacenters(&self) -> &RwLock<Vec<Datacenter>> {
383        &self.datacenters
384    }
385
386    /// Borrow the per-peer auto-eject deciders.
387    #[must_use]
388    pub fn auto_eject(&self) -> &RwLock<Vec<AutoEject>> {
389        &self.auto_eject
390    }
391
392    /// Rebuild the per-rack token continuum from the current peer
393    /// table. Mirrors `vnode_update`. When the configured
394    /// distribution is
395    /// [`crate::conf::Distribution::RandomSlicing`], a
396    /// [`crate::hashkit::random_slicing::RandomSlices`] table is
397    /// installed on each rack alongside the continuum so the
398    /// shadow-distribution path can still walk the vnode view.
399    pub fn rebuild_ring(&self) {
400        let peers = self.peers.read();
401        let mut dcs = self.datacenters.write();
402        // Make sure all (dc, rack) pairs exist.
403        for p in peers.iter() {
404            let dc_idx = if let Some(i) = dcs.iter().position(|d| d.name() == p.dc()) {
405                i
406            } else {
407                dcs.push(Datacenter::new(p.dc().to_string()));
408                dcs.len() - 1
409            };
410            dcs[dc_idx].upsert_rack(p.rack().to_string());
411        }
412        let entries: Vec<_> = peers
413            .iter()
414            .map(|p| crate::cluster::vnode::PeerTokens {
415                peer_idx: p.idx(),
416                dc: p.dc(),
417                rack: p.rack(),
418                tokens: p.tokens(),
419            })
420            .collect();
421        crate::cluster::vnode::rebuild_continuums(&mut dcs, &entries);
422        let live_or_shadow_uses_rs = matches!(
423            self.config.distribution,
424            crate::conf::Distribution::RandomSlicing
425        ) || matches!(
426            self.config.distribution_shadow,
427            Some(crate::conf::Distribution::RandomSlicing)
428        );
429        if live_or_shadow_uses_rs {
430            Self::install_random_slices(&peers, &mut dcs);
431        }
432    }
433
434    fn install_random_slices(peers: &[crate::cluster::peer::Peer], dcs: &mut [Datacenter]) {
435        use crate::hashkit::random_slicing::RandomSlices;
436        for dc in dcs.iter_mut() {
437            let dc_name = dc.name().to_string();
438            for rack in dc.racks_mut().iter_mut() {
439                let mut names: Vec<String> = peers
440                    .iter()
441                    .filter(|p| p.dc() == dc_name && p.rack() == rack.name())
442                    .map(|p| p.endpoint().pname())
443                    .collect();
444                names.sort();
445                names.dedup();
446                if names.is_empty() {
447                    continue;
448                }
449                let refs: Vec<&str> = names.iter().map(String::as_str).collect();
450                if let Ok(slices) = RandomSlices::from_uniform(&refs) {
451                    rack.set_random_slices(slices);
452                } else {
453                    tracing::warn!(
454                        target: "dynomite::cluster::pool",
455                        rack = rack.name(),
456                        dc = %dc_name,
457                        "random-slicing build failed; falling back to vnode for this rack"
458                    );
459                }
460            }
461        }
462    }
463
464    /// Walk the datacenters and choose, for each remote DC, a rack
465    /// for cross-DC replication. Mirrors
466    /// `preselect_remote_rack_for_replication`.
467    pub fn preselect_remote_racks(&self) {
468        let mut dcs = self.datacenters.write();
469        for dc in dcs.iter_mut() {
470            dc.sort_racks();
471        }
472        // Find the index of the local rack in the local DC.
473        let mut my_rack_index = 0usize;
474        for dc in dcs.iter() {
475            if dc.name() == self.config.dc {
476                if let Some(idx) = dc.rack_idx(&self.config.rack) {
477                    my_rack_index = idx;
478                }
479                break;
480            }
481        }
482        for dc in dcs.iter_mut() {
483            if dc.name() == self.config.dc {
484                dc.set_preselected_rack_idx(None);
485                continue;
486            }
487            let rack_count = dc.racks().len();
488            if rack_count == 0 {
489                dc.set_preselected_rack_idx(None);
490            } else {
491                dc.set_preselected_rack_idx(Some(my_rack_index % rack_count));
492            }
493        }
494    }
495
496    /// Initialise a per-DC [`crate::msg::ResponseMgr`] vector for
497    /// the supplied request. The walker visits every datacenter
498    /// and produces one manager per DC sized to the rack count.
499    /// Mirrors `init_response_mgr_all_dcs`.
500    ///
501    /// # Examples
502    ///
503    /// ```
504    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
505    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
506    /// # use dynomite::hashkit::DynToken;
507    /// # use dynomite::conf::{DataStore, HashType};
508    /// # use dynomite::msg::{ConsistencyLevel, Msg, MsgType};
509    /// # let cfg = PoolConfig::default();
510    /// # let local = Peer::new(
511    /// #    0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
512    /// #    vec![DynToken::from_u32(0)], true, true, false,
513    /// # );
514    /// let pool = ServerPool::new(cfg, vec![local]);
515    /// let req = Msg::new(1, MsgType::ReqRedisGet, true);
516    /// let mgrs = pool.init_response_mgrs(&req);
517    /// assert_eq!(mgrs.len(), 1);
518    /// ```
519    #[must_use]
520    pub fn init_response_mgrs(&self, req: &crate::msg::Msg) -> Vec<crate::msg::ResponseMgr> {
521        use crate::msg::{ResponseMgr, MAX_REPLICAS_PER_DC};
522        let dcs = self.datacenters.read();
523        let mut out = Vec::with_capacity(dcs.len());
524        for dc in dcs.iter() {
525            let rack_count = dc.racks().len();
526            let max_responses = u8::try_from(rack_count.clamp(1, MAX_REPLICAS_PER_DC))
527                .unwrap_or(u8::try_from(MAX_REPLICAS_PER_DC).unwrap_or(1));
528            out.push(ResponseMgr::new(
529                req,
530                max_responses,
531                Some(dc.name().to_string()),
532            ));
533        }
534        out
535    }
536}
537
538#[cfg(test)]
539mod tests {
540    use super::*;
541    use crate::cluster::peer::PeerEndpoint;
542    use crate::hashkit::DynToken;
543
544    fn cfg(dc: &str, rack: &str) -> PoolConfig {
545        PoolConfig {
546            dc: dc.into(),
547            rack: rack.into(),
548            ..PoolConfig::default()
549        }
550    }
551
552    fn peer(idx: u32, dc: &str, rack: &str, tok: u32, is_local: bool, is_same: bool) -> Peer {
553        Peer::new(
554            idx,
555            PeerEndpoint::tcp("127.0.0.1".into(), 8101 + u16::try_from(idx).unwrap_or(0)),
556            rack.into(),
557            dc.into(),
558            vec![DynToken::from_u32(tok)],
559            is_local,
560            is_same,
561            false,
562        )
563    }
564
565    #[test]
566    fn build_pool_populates_topology() {
567        let pool = ServerPool::new(
568            cfg("dc1", "r1"),
569            vec![
570                peer(0, "dc1", "r1", 10, true, true),
571                peer(1, "dc1", "r2", 20, false, true),
572                peer(2, "dc2", "r1", 30, false, false),
573            ],
574        );
575        let topology = pool.datacenters().read();
576        let dc1 = topology.iter().find(|d| d.name() == "dc1").unwrap();
577        assert_eq!(dc1.racks().len(), 2);
578    }
579
580    #[test]
581    fn preselect_remote_picks_per_dc() {
582        let pool = ServerPool::new(
583            cfg("dc1", "rA"),
584            vec![
585                peer(0, "dc1", "rA", 10, true, true),
586                peer(1, "dc2", "rA", 20, false, false),
587                peer(2, "dc2", "rB", 30, false, false),
588            ],
589        );
590        pool.preselect_remote_racks();
591        let topology = pool.datacenters().read();
592        let dc2 = topology.iter().find(|d| d.name() == "dc2").unwrap();
593        // Local rack "rA" is at sorted index 0, dc2 has 2 racks, so
594        // preselected idx is 0 -> "rA".
595        assert_eq!(
596            dc2.preselected_rack()
597                .map(super::super::datacenter::Rack::name),
598            Some("rA")
599        );
600    }
601
602    #[test]
603    fn init_response_mgrs_one_per_dc() {
604        let pool = ServerPool::new(
605            cfg("dc1", "r1"),
606            vec![
607                peer(0, "dc1", "r1", 10, true, true),
608                peer(1, "dc2", "r1", 20, false, false),
609            ],
610        );
611        let req = crate::msg::Msg::new(1, crate::msg::MsgType::ReqRedisGet, true);
612        let mgrs = pool.init_response_mgrs(&req);
613        assert_eq!(mgrs.len(), 2);
614    }
615}