Skip to main content

dynomite/cluster/
dispatch.rs

1//! Cluster-aware [`Dispatcher`].
2//!
3//! Routes parsed [`Msg`]s based on the configured consistency level
4//! and the [`crate::cluster::pool::ServerPool`] topology:
5//!
6//! * `DC_ONE` reads pick the rack-local replica via the snitch.
7//! * `DC_ONE` writes fan out to every replica in the local DC.
8//! * `DC_QUORUM` / `DC_SAFE_QUORUM` reads fan out to every replica
9//!   in the local DC.
10//! * `DC_EACH_SAFE_QUORUM` writes fan out per-DC, walking the
11//!   per-DC racks via the preselected rack from
12//!   [`crate::cluster::pool::ServerPool::preselect_remote_racks`].
13//!
14//! The actual outbound delivery happens through the per-peer
15//! [`crate::net::ConnPool`]s; this module produces a
16//! [`DispatchPlan`] (the list of replica peers a request must be
17//! routed to) and exposes the planning logic so it can be tested
18//! independently of the runtime fan-out.
19//!
20//! # Examples
21//!
22//! ```
23//! use dynomite::cluster::dispatch::{ClusterDispatcher, DispatchPlan};
24//! use dynomite::cluster::pool::{PoolConfig, ServerPool};
25//! use dynomite::cluster::peer::{Peer, PeerEndpoint};
26//! use dynomite::hashkit::DynToken;
27//! use dynomite::msg::{Msg, MsgType};
28//! use std::sync::Arc;
29//!
30//! let cfg = PoolConfig {
31//!     dc: "d".into(), rack: "r".into(),
32//!     ..PoolConfig::default()
33//! };
34//! let local = Peer::new(
35//!     0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
36//!     vec![DynToken::from_u32(0)], true, true, false,
37//! );
38//! let pool = Arc::new(ServerPool::new(cfg, vec![local]));
39//! let disp = ClusterDispatcher::new(pool);
40//! let req = Msg::new(1, MsgType::ReqRedisGet, true);
41//! let plan = disp.plan(&req, b"foo");
42//! assert!(matches!(plan, DispatchPlan::LocalDatastore));
43//! ```
44
45use std::sync::atomic::{AtomicU64, Ordering};
46use std::sync::Arc;
47
48use tokio::sync::mpsc;
49
50use crate::cluster::pool::ServerPool;
51use crate::cluster::snitch::{rack_distance, RackDistance};
52use crate::cluster::vnode;
53use crate::conf::HashType as ConfHashType;
54use crate::hashkit::{self, HashType};
55use crate::io::mbuf::MbufPool;
56use crate::msg::{ConsistencyLevel, Msg, MsgRouting, MsgType};
57use crate::net::dispatcher::{DispatchOutcome, Dispatcher, OutboundEnvelope, ServerSink};
58use crate::net::server::OutboundRequest;
59
60/// Process-global Prometheus-friendly counter of
61/// shadow-distribution disagreements. The dispatcher bumps this
62/// counter every time the configured `distribution` and the
63/// configured `distribution_shadow` choose different peers for
64/// the same key. Exposed for both the stats endpoint and the
65/// integration test that exercises shadow mode.
66///
67/// # Examples
68///
69/// ```
70/// use dynomite::cluster::dispatch::distribution_shadow_disagreement_total;
71/// let _seen = distribution_shadow_disagreement_total();
72/// ```
73#[must_use]
74pub fn distribution_shadow_disagreement_total() -> u64 {
75    SHADOW_DISAGREEMENTS.load(Ordering::Relaxed)
76}
77
78/// Reset the shadow-disagreement counter. Used by integration
79/// tests that need a clean baseline; never called from
80/// production code.
81///
82/// # Examples
83///
84/// ```
85/// use dynomite::cluster::dispatch::reset_distribution_shadow_disagreement_total;
86/// reset_distribution_shadow_disagreement_total();
87/// ```
88pub fn reset_distribution_shadow_disagreement_total() {
89    SHADOW_DISAGREEMENTS.store(0, Ordering::Relaxed);
90}
91
92static SHADOW_DISAGREEMENTS: AtomicU64 = AtomicU64::new(0);
93
94fn bump_shadow_disagreement() {
95    SHADOW_DISAGREEMENTS.fetch_add(1, Ordering::Relaxed);
96}
97
98/// Build the `dispatch.plan` info span and enter it. Returns the
99/// originating client request span (captured before the plan
100/// span was entered) plus the entered plan-span guard. Factored
101/// out so [`ClusterDispatcher::dispatch`] stays inside the
102/// project's per-function line budget.
103fn enter_plan_span(
104    req_id: u64,
105    plan: &DispatchPlan,
106) -> (tracing::Span, tracing::span::EnteredSpan) {
107    let req_span = tracing::Span::current();
108    let kind: &'static str = match plan {
109        DispatchPlan::Drop => "drop",
110        DispatchPlan::NoTargets => "no_targets",
111        DispatchPlan::LocalDatastore => "local_datastore",
112        DispatchPlan::Replicas { .. } => "replicas",
113    };
114    let targets = match plan {
115        DispatchPlan::Replicas { targets, .. } => targets.len(),
116        _ => 0,
117    };
118    let span = tracing::info_span!("dispatch.plan", req_id, plan = kind, targets,).entered();
119    (req_span, span)
120}
121
122/// Map a configuration-layer [`ConfHashType`] to the runtime
123/// [`HashType`] used by the ring hash. The two enums are distinct
124/// (one is parsed from YAML, the other drives
125/// [`crate::hashkit::hash64`]); this is the single conversion seam
126/// so the dispatcher, the reaper, and the dyniak replica router all
127/// agree on which hash a pool uses.
128#[must_use]
129pub fn map_hash(h: ConfHashType) -> HashType {
130    match h {
131        ConfHashType::OneAtATime => HashType::OneAtATime,
132        ConfHashType::Md5 => HashType::Md5,
133        ConfHashType::Crc16 => HashType::Crc16,
134        ConfHashType::Crc32 => HashType::Crc32,
135        ConfHashType::Crc32a => HashType::Crc32a,
136        ConfHashType::Fnv1_64 => HashType::Fnv1_64,
137        ConfHashType::Fnv1a64 => HashType::Fnv1a_64,
138        ConfHashType::Fnv1_32 => HashType::Fnv1_32,
139        ConfHashType::Fnv1a32 => HashType::Fnv1a_32,
140        ConfHashType::Hsieh => HashType::Hsieh,
141        ConfHashType::Murmur => HashType::Murmur,
142        ConfHashType::Jenkins => HashType::Jenkins,
143        ConfHashType::Murmur3 => HashType::Murmur3,
144        ConfHashType::Murmur3X64_64 => HashType::Murmur3X64_64,
145    }
146}
147
148/// One replica target produced by [`ClusterDispatcher::plan`].
149#[derive(Clone, Debug, Eq, PartialEq)]
150pub struct ReplicaTarget {
151    /// Index of the target peer in the pool's peer array.
152    pub peer_idx: u32,
153    /// Datacenter name.
154    pub dc: String,
155    /// Rack name.
156    pub rack: String,
157    /// True when the target is the local node.
158    pub is_local: bool,
159}
160
161/// Dispatch plan produced by the cluster dispatcher.
162///
163/// `LocalDatastore` is the early-return branch the reference
164/// engine takes when the routing tag is `ROUTING_LOCAL_NODE_ONLY`
165/// (or when the request is destined for the local node and the
166/// topology has only one peer); the per-connection driver then
167/// hands the request off to its server-side connection pool.
168#[derive(Clone, Debug, Eq, PartialEq)]
169pub enum DispatchPlan {
170    /// Hand the request straight to the local datastore.
171    LocalDatastore,
172    /// Forward to one or more peer replicas. The carried
173    /// consistency level is the one the planner resolved for
174    /// this request (after applying any bucket-type override),
175    /// so the dispatcher's reply coalescer does not have to
176    /// re-resolve it.
177    Replicas {
178        /// Replica peers the request must be routed to.
179        targets: Vec<ReplicaTarget>,
180        /// Resolved consistency level.
181        consistency: ConsistencyLevel,
182    },
183    /// Reply with an error: the cluster has no quorum-eligible
184    /// targets.
185    NoTargets,
186    /// Drop the request (`QUIT`-style swallow).
187    Drop,
188}
189
190/// Cluster-aware dispatcher.
191#[derive(Clone)]
192pub struct ClusterDispatcher {
193    pool: Arc<ServerPool>,
194    /// Outbound channel feeding the local datastore driver. When
195    /// `None`, `LocalDatastore` plans short-circuit to `Pending`
196    /// without forwarding (used by tests that do not need a real
197    /// backend). When set, requests for the local node are
198    /// encoded onto the wire and shipped to the [`crate::net::ServerConn`]
199    /// task that drives the redis / memcache backend.
200    backend: Option<mpsc::Sender<OutboundRequest>>,
201    /// Per-peer outbound channel for cross-DC fan-out. Keyed by
202    /// `Peer::idx`. When a `DispatchPlan::Replicas` plan names a
203    /// non-local peer, the dispatcher forwards via the matching
204    /// channel to a `DnodeServerConn` task. Peers without a
205    /// wired channel are skipped (`Pending`); when no replica
206    /// is reachable for the consistency level the dispatcher
207    /// falls back to a `DynomiteNoQuorumAchieved` error response.
208    peer_backends: std::collections::HashMap<u32, mpsc::Sender<OutboundRequest>>,
209    /// Mbuf pool used to render synthetic error payloads.
210    /// `MbufPool` already wraps an `Arc`, so cloning the
211    /// dispatcher (and the pool with it) shares the same free
212    /// list across every cluster handle.
213    mbuf_pool: MbufPool,
214    /// Optional node-local hint store. When set AND the pool's
215    /// `enable_hinted_handoff` flag is true, write requests
216    /// targeted at peers in [`crate::cluster::peer::PeerState::Down`]
217    /// (or at peers whose outbound channel is closed / full) are
218    /// recorded as hints and counted toward the consistency
219    /// threshold, instead of being silently skipped.
220    hint_store: Option<Arc<crate::cluster::hints::HintStore>>,
221    /// Optional failure-cause metrics handle. When wired, every
222    /// error-producing branch in the dispatcher increments the
223    /// matching counter via the [`crate::stats::FailureMetrics`]
224    /// accumulator. When `None`, the dispatcher's behaviour is
225    /// unchanged.
226    failure_metrics: Option<Arc<crate::stats::FailureMetrics>>,
227    /// Optional command-dispatch extension. When set, the
228    /// dispatcher offers FT.* / HSET requests to the extension
229    /// before the routing planner runs; the extension may
230    /// short-circuit with a synthesised reply
231    /// ([`crate::net::DispatchOutcome::Inline`]), reject the
232    /// request with a structured error
233    /// ([`crate::net::DispatchOutcome::Error`]), or fall
234    /// through to the standard storage path. When `None`, the
235    /// dispatcher's behaviour is unchanged. See
236    /// [`crate::embed::CommandExtension`] for the trait shape.
237    command_extension: Option<Arc<dyn crate::embed::CommandExtension>>,
238    /// Optional in-process local datastore hook. When set, a
239    /// [`DispatchPlan::LocalDatastore`] request is handed to this
240    /// [`crate::embed::Datastore`] (the parsed request `Msg`,
241    /// asynchronously) instead of being relayed over the
242    /// [`Self::backend`] byte channel. The embedded server wires
243    /// this so a connection accepted on its `listen:` socket is
244    /// served through the same `Datastore` hook that
245    /// `ServerHandle::inject_request` uses. When `None`, the
246    /// backend byte channel is used (the standalone proxy path).
247    local_datastore: Option<Arc<dyn crate::embed::hooks::Datastore>>,
248}
249
250impl std::fmt::Debug for ClusterDispatcher {
251    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
252        // `Datastore` is not `Debug` (it is a public embedding
253        // trait and must stay object-safe without a Debug bound),
254        // so the local-datastore hook is reported as a presence
255        // flag rather than its contents.
256        f.debug_struct("ClusterDispatcher")
257            .field("backend", &self.backend.is_some())
258            .field("peer_backends", &self.peer_backends.len())
259            .field("hint_store", &self.hint_store.is_some())
260            .field("failure_metrics", &self.failure_metrics.is_some())
261            .field("command_extension", &self.command_extension)
262            .field("local_datastore", &self.local_datastore.is_some())
263            .finish_non_exhaustive()
264    }
265}
266
267impl ClusterDispatcher {
268    /// Wrap a [`ServerPool`] in a dispatcher.
269    ///
270    /// # Examples
271    ///
272    /// ```
273    /// # use dynomite::cluster::dispatch::ClusterDispatcher;
274    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
275    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
276    /// # use dynomite::hashkit::DynToken;
277    /// # use std::sync::Arc;
278    /// # let cfg = PoolConfig::default();
279    /// # let local = Peer::new(
280    /// #    0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
281    /// #    vec![DynToken::from_u32(0)], true, true, false,
282    /// # );
283    /// let pool = Arc::new(ServerPool::new(cfg, vec![local]));
284    /// let disp = ClusterDispatcher::new(pool);
285    /// let _ = disp.pool();
286    /// ```
287    #[must_use]
288    pub fn new(pool: Arc<ServerPool>) -> Self {
289        Self {
290            pool,
291            backend: None,
292            peer_backends: std::collections::HashMap::new(),
293            mbuf_pool: MbufPool::default(),
294            hint_store: None,
295            failure_metrics: None,
296            command_extension: None,
297            local_datastore: None,
298        }
299    }
300
301    /// Override the dispatcher's mbuf pool. Useful when the
302    /// embedding wants every synthetic error payload to come from
303    /// the same recycled buffers as the rest of the engine.
304    ///
305    /// # Examples
306    ///
307    /// ```
308    /// # use dynomite::cluster::dispatch::ClusterDispatcher;
309    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
310    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
311    /// # use dynomite::hashkit::DynToken;
312    /// # use dynomite::io::mbuf::MbufPool;
313    /// # use std::sync::Arc;
314    /// # let cfg = PoolConfig::default();
315    /// # let local = Peer::new(
316    /// #    0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
317    /// #    vec![DynToken::from_u32(0)], true, true, false,
318    /// # );
319    /// let pool = Arc::new(ServerPool::new(cfg, vec![local]));
320    /// let _disp = ClusterDispatcher::new(pool).with_mbuf_pool(MbufPool::default());
321    /// ```
322    #[must_use]
323    pub fn with_mbuf_pool(mut self, pool: MbufPool) -> Self {
324        self.mbuf_pool = pool;
325        self
326    }
327
328    /// Borrow the dispatcher's mbuf pool. Exposed so embedders
329    /// can reuse the same pool when building synthetic responses
330    /// outside the dispatcher's own code paths.
331    #[must_use]
332    pub fn mbuf_pool(&self) -> &MbufPool {
333        &self.mbuf_pool
334    }
335
336    /// Attach a backend request channel. Calls to [`Self::dispatch`]
337    /// that produce a [`DispatchPlan::LocalDatastore`] plan will
338    /// forward the request bytes onto this channel for the local
339    /// datastore driver to write to the backend.
340    ///
341    /// The channel sender must be the request side of a
342    /// [`crate::net::ServerConn`] task; multiple senders cloned from
343    /// the same channel are fine.
344    #[must_use]
345    pub fn with_backend(mut self, backend: mpsc::Sender<OutboundRequest>) -> Self {
346        self.backend = Some(backend);
347        self
348    }
349
350    /// Attach an in-process local datastore hook. A
351    /// [`DispatchPlan::LocalDatastore`] request is then handed to
352    /// the supplied [`crate::embed::hooks::Datastore`] (the parsed
353    /// request `Msg`) and its reply is delivered on the
354    /// connection's responder, instead of being relayed over the
355    /// [`Self::with_backend`] byte channel. The embedded server
356    /// wires this so a client accepted on its `listen:` socket is
357    /// served through the same `Datastore` hook that
358    /// `ServerHandle::inject_request` uses. Takes precedence over
359    /// the backend channel for local requests.
360    #[must_use]
361    pub fn with_local_datastore(
362        mut self,
363        datastore: Arc<dyn crate::embed::hooks::Datastore>,
364    ) -> Self {
365        self.local_datastore = Some(datastore);
366        self
367    }
368
369    /// Attach an outbound channel for a single peer (by
370    /// `Peer::idx`). The supplied sender feeds a
371    /// [`crate::net::DnodeServerConn`] task that writes
372    /// dnode-framed requests to the peer's `dyn_listen` and
373    /// routes the response back through the per-request
374    /// responder channel.
375    ///
376    /// Wiring is additive: call this once per non-local peer.
377    /// Calling it again with the same `peer_idx` replaces the
378    /// previous sender (used by reconnect supervisors that
379    /// rebuild channels on restart).
380    #[must_use]
381    pub fn with_peer_backend(
382        mut self,
383        peer_idx: u32,
384        sender: mpsc::Sender<OutboundRequest>,
385    ) -> Self {
386        self.peer_backends.insert(peer_idx, sender);
387        self
388    }
389
390    /// Whether a backend channel is wired.
391    #[must_use]
392    pub fn has_backend(&self) -> bool {
393        self.backend.is_some()
394    }
395
396    /// Number of peer-backend channels wired.
397    #[must_use]
398    pub fn peer_backend_count(&self) -> usize {
399        self.peer_backends.len()
400    }
401
402    /// Borrow the underlying pool.
403    #[must_use]
404    pub fn pool(&self) -> &Arc<ServerPool> {
405        &self.pool
406    }
407
408    /// Attach a [`crate::cluster::hints::HintStore`].
409    ///
410    /// When set AND the pool's `enable_hinted_handoff` flag is
411    /// `true`, write requests targeted at peers in
412    /// [`crate::cluster::peer::PeerState::Down`] (or at peers
413    /// whose outbound channel is closed / full) are stored as
414    /// hints and counted toward the consistency threshold. The
415    /// background drainer task in `dynomited` is responsible
416    /// for shipping the hints back to the peer once it returns
417    /// to [`crate::cluster::peer::PeerState::Normal`]. Without
418    /// this builder call (or with `enable_hinted_handoff: false`)
419    /// the dispatcher behaviour is unchanged: a Down or
420    /// unreachable target is silently skipped.
421    ///
422    /// # Examples
423    ///
424    /// ```
425    /// # use std::sync::Arc;
426    /// # use dynomite::cluster::dispatch::ClusterDispatcher;
427    /// # use dynomite::cluster::hints::HintStore;
428    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
429    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
430    /// # use dynomite::hashkit::DynToken;
431    /// let cfg = PoolConfig { enable_hinted_handoff: true, ..PoolConfig::default() };
432    /// let local = Peer::new(
433    ///     0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
434    ///     vec![DynToken::from_u32(0)], true, true, false,
435    /// );
436    /// let pool = Arc::new(ServerPool::new(cfg, vec![local]));
437    /// let store = Arc::new(HintStore::new(64 * 1024 * 1024));
438    /// let _disp = ClusterDispatcher::new(pool).with_hint_store(store);
439    /// ```
440    #[must_use]
441    pub fn with_hint_store(mut self, store: Arc<crate::cluster::hints::HintStore>) -> Self {
442        self.hint_store = Some(store);
443        self
444    }
445
446    /// Borrow the wired hint store, if any.
447    #[must_use]
448    pub fn hint_store(&self) -> Option<&Arc<crate::cluster::hints::HintStore>> {
449        self.hint_store.as_ref()
450    }
451
452    /// Attach a [`crate::stats::FailureMetrics`] handle.
453    ///
454    /// When wired, each error-producing branch in the
455    /// dispatcher (no-targets, peer-channel-full,
456    /// peer-channel-closed, backend-channel-full,
457    /// backend-channel-closed, response-timeout) increments
458    /// the matching counter so an operator can pull the
459    /// per-cause histogram off the `/stats` and `/metrics`
460    /// endpoints. The default behaviour is unchanged when no
461    /// metrics handle is supplied.
462    ///
463    /// # Examples
464    ///
465    /// ```
466    /// # use std::sync::Arc;
467    /// # use dynomite::cluster::dispatch::ClusterDispatcher;
468    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
469    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
470    /// # use dynomite::hashkit::DynToken;
471    /// # use dynomite::stats::FailureMetrics;
472    /// let cfg = PoolConfig::default();
473    /// let local = Peer::new(
474    ///     0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
475    ///     vec![DynToken::from_u32(0)], true, true, false,
476    /// );
477    /// let pool = Arc::new(ServerPool::new(cfg, vec![local]));
478    /// let metrics = Arc::new(FailureMetrics::new());
479    /// let _disp = ClusterDispatcher::new(pool).with_failure_metrics(metrics);
480    /// ```
481    #[must_use]
482    pub fn with_failure_metrics(mut self, metrics: Arc<crate::stats::FailureMetrics>) -> Self {
483        self.failure_metrics = Some(metrics);
484        self
485    }
486
487    /// Borrow the wired failure-metrics handle, if any.
488    #[must_use]
489    pub fn failure_metrics(&self) -> Option<&Arc<crate::stats::FailureMetrics>> {
490        self.failure_metrics.as_ref()
491    }
492
493    /// Attach a [`crate::embed::CommandExtension`].
494    ///
495    /// When wired, parsed FT.* requests and HSET requests are
496    /// offered to the extension before the routing planner
497    /// runs. The extension may produce a synthesised RESP
498    /// reply (returned to the client as
499    /// [`crate::net::DispatchOutcome::Inline`]), surface a
500    /// structured `-ERR ...` reply
501    /// ([`crate::net::DispatchOutcome::Error`]), or fall
502    /// through. Without this builder call the dispatcher's
503    /// behaviour is unchanged: FT.* keywords are forwarded to
504    /// the local datastore (which typically rejects them with
505    /// `-ERR unknown command`) and HSETs proceed unmodified.
506    ///
507    /// The standard RediSearch implementation lives in the
508    /// `dynomite-search` crate; the trait itself is part of
509    /// the engine so embedders can plug their own.
510    ///
511    /// # Examples
512    ///
513    /// ```
514    /// # use std::sync::Arc;
515    /// # use dynomite::cluster::dispatch::ClusterDispatcher;
516    /// # use dynomite::cluster::peer::{Peer, PeerEndpoint};
517    /// # use dynomite::cluster::pool::{PoolConfig, ServerPool};
518    /// # use dynomite::embed::{CommandExtension, HsetOutcome};
519    /// # use dynomite::hashkit::DynToken;
520    /// # use dynomite::msg::MsgType;
521    /// #[derive(Debug)]
522    /// struct NoOp;
523    /// impl CommandExtension for NoOp {
524    ///     fn handles_msg_type(&self, _: MsgType) -> bool { false }
525    ///     fn try_dispatch(&self, _: &[&[u8]]) -> Option<Vec<u8>> { None }
526    /// }
527    /// let cfg = PoolConfig::default();
528    /// let local = Peer::new(
529    ///     0, PeerEndpoint::tcp("h".into(), 1), "r".into(), "d".into(),
530    ///     vec![DynToken::from_u32(0)], true, true, false,
531    /// );
532    /// let pool = Arc::new(ServerPool::new(cfg, vec![local]));
533    /// let _disp = ClusterDispatcher::new(pool).with_command_extension(Arc::new(NoOp));
534    /// ```
535    #[must_use]
536    pub fn with_command_extension(mut self, ext: Arc<dyn crate::embed::CommandExtension>) -> Self {
537        self.command_extension = Some(ext);
538        self
539    }
540
541    /// Borrow the wired command extension, if any.
542    #[must_use]
543    pub fn command_extension(&self) -> Option<&Arc<dyn crate::embed::CommandExtension>> {
544        self.command_extension.as_ref()
545    }
546
547    /// True when both the hint store is wired AND the pool has
548    /// `enable_hinted_handoff: true`. Hot-path predicate.
549    #[must_use]
550    pub fn hinted_handoff_active(&self) -> bool {
551        self.hint_store.is_some() && self.pool.config().enable_hinted_handoff
552    }
553
554    /// Compute the routing plan for `req` with the supplied key.
555    ///
556    /// `key` is the primary key of the request (the first key
557    /// returned by [`Msg::keys`] for parsed redis / memcache
558    /// commands, or an empty slice for argument-less commands).
559    ///
560    /// The function never panics; it consults the live peer table
561    /// behind the pool's `RwLock` and returns
562    /// [`DispatchPlan::NoTargets`] when the topology cannot
563    /// satisfy the request.
564    ///
565    /// # Examples
566    ///
567    /// See the module-level example.
568    #[must_use]
569    pub fn plan(&self, req: &Msg, key: &[u8]) -> DispatchPlan {
570        let cfg = self.pool.config();
571        let peers = self.pool.peers().read();
572        if peers.is_empty() {
573            self.record_no_targets_metric(cfg, ConsistencyLevel::default());
574            return DispatchPlan::NoTargets;
575        }
576        if matches!(req.routing(), MsgRouting::LocalNodeOnly) {
577            return DispatchPlan::LocalDatastore;
578        }
579        if key.is_empty() {
580            return DispatchPlan::LocalDatastore;
581        }
582        let token = hashkit::hash(map_hash(cfg.hash), key);
583        let key_hash64 = hashkit::hash64(map_hash(cfg.hash), key);
584        let bucket = crate::proto::redis::bucket_name(key);
585        let bucket_type = cfg.resolve_bucket_type(bucket);
586        let is_read = matches!(req.ty(), MsgType::Unknown) || req.flags().is_read;
587        let consistency = match (bucket_type, is_read) {
588            (Some(bt), true) => bt.read_consistency,
589            (Some(bt), false) => bt.write_consistency,
590            (None, true) => cfg.read_consistency,
591            (None, false) => cfg.write_consistency,
592        };
593        let n_val_cap = bucket_type.map_or(0, |bt| bt.n_val);
594        let dcs = self.pool.datacenters().read();
595        // When hinted handoff is active and the request is a
596        // write, peers in `Down` are kept in the routable set
597        // so the dispatcher can hint them at fan-out time. The
598        // hint counts toward the consistency threshold; the
599        // background drainer ships the hint once the peer
600        // returns to `Normal`. For reads (and for writes when
601        // handoff is off), Down peers are filtered out as
602        // before.
603        let include_down = self.hinted_handoff_active() && !is_read;
604        let routable = collect_routable(
605            &dcs,
606            &peers,
607            &token,
608            key_hash64,
609            cfg.distribution,
610            include_down,
611        );
612        if let Some(shadow) = cfg.distribution_shadow {
613            if shadow != cfg.distribution {
614                let shadow_routable =
615                    collect_routable(&dcs, &peers, &token, key_hash64, shadow, include_down);
616                if !plans_agree(&routable, &shadow_routable) {
617                    bump_shadow_disagreement();
618                    tracing::debug!(
619                        target: "dynomite::dispatch::shadow",
620                        live = cfg.distribution.as_str(),
621                        shadow = shadow.as_str(),
622                        "shadow distribution disagreed on key route"
623                    );
624                }
625            }
626        }
627        if routable.is_empty() {
628            self.record_no_targets_metric(cfg, consistency);
629            return DispatchPlan::NoTargets;
630        }
631        let (local, remote): (Vec<_>, Vec<_>) = routable
632            .into_iter()
633            .partition(|(dc_idx, _, _)| dcs[*dc_idx].name() == cfg.dc);
634        let plan = plan_with_consistency(
635            cfg,
636            &dcs,
637            &peers,
638            consistency,
639            req.routing(),
640            is_read,
641            RoutablePartition { local, remote },
642        );
643        let plan = cap_replicas(plan, n_val_cap);
644        if matches!(plan, DispatchPlan::NoTargets) {
645            self.record_no_targets_metric(cfg, consistency);
646        }
647        plan
648    }
649
650    /// Record a `dispatch_no_targets_total` metric tick using
651    /// the local-DC labels, when a metrics handle is wired.
652    fn record_no_targets_metric(
653        &self,
654        cfg: &crate::cluster::pool::PoolConfig,
655        consistency: ConsistencyLevel,
656    ) {
657        if let Some(m) = self.failure_metrics.as_ref() {
658            m.record_no_targets(&cfg.dc, &cfg.rack, consistency);
659        }
660    }
661
662    /// Resolve the destination DC of a peer (for per-peer
663    /// failure metrics). Local peers and unknown indexes both
664    /// fall back to the configured local DC.
665    fn peer_dc_label(&self, peer_idx: u32) -> String {
666        let peers = self.pool.peers().read();
667        peers
668            .get(peer_idx as usize)
669            .map_or_else(|| self.pool.config().dc.clone(), |p| p.dc().to_string())
670    }
671}
672
673/// Apply the bucket-type `n_val` fan-out cap to a freshly
674/// computed plan. Only `DispatchPlan::Replicas` is affected; the
675/// other variants pass through unchanged. `cap == 0` means "no
676/// cap" and is the no-op used for keys without a matching bucket
677/// type.
678fn cap_replicas(plan: DispatchPlan, cap: u8) -> DispatchPlan {
679    if cap == 0 {
680        return plan;
681    }
682    let cap = cap as usize;
683    match plan {
684        DispatchPlan::Replicas {
685            mut targets,
686            consistency,
687        } if targets.len() > cap => {
688            targets.truncate(cap);
689            DispatchPlan::Replicas {
690                targets,
691                consistency,
692            }
693        }
694        other => other,
695    }
696}
697
698fn plans_agree(a: &[(usize, usize, u32)], b: &[(usize, usize, u32)]) -> bool {
699    if a.len() != b.len() {
700        return false;
701    }
702    let mut a_idx: Vec<u32> = a.iter().map(|t| t.2).collect();
703    let mut b_idx: Vec<u32> = b.iter().map(|t| t.2).collect();
704    a_idx.sort_unstable();
705    b_idx.sort_unstable();
706    a_idx == b_idx
707}
708
709fn collect_routable(
710    dcs: &[crate::cluster::Datacenter],
711    peers: &[crate::cluster::peer::Peer],
712    token: &crate::hashkit::DynToken,
713    hash64: u64,
714    distribution: crate::conf::Distribution,
715    include_down: bool,
716) -> Vec<(usize, usize, u32)> {
717    let mut routable: Vec<(usize, usize, u32)> = Vec::new();
718    for (dc_idx, dc) in dcs.iter().enumerate() {
719        for (rack_idx, rack) in dc.racks().iter().enumerate() {
720            let candidate = match (distribution, rack.random_slices()) {
721                (crate::conf::Distribution::RandomSlicing, Some(slices)) => {
722                    // Map the chosen claimant name back onto a
723                    // peer index. The slice table holds peer
724                    // pname strings (host:port) so the
725                    // resolution is a linear scan over the
726                    // rack's peer set, which is small (peers
727                    // per rack, not the whole pool).
728                    slices.claimant_for(hash64).and_then(|name| {
729                        peers.iter().find_map(|p| {
730                            if p.dc() == dc.name()
731                                && p.rack() == rack.name()
732                                && p.endpoint().pname() == name
733                            {
734                                Some(p.idx())
735                            } else {
736                                None
737                            }
738                        })
739                    })
740                }
741                _ => vnode::dispatch(rack.continuums(), token),
742            };
743            if let Some(peer_idx) = candidate {
744                if let Some(peer) = peers.get(peer_idx as usize) {
745                    let state = peer.state();
746                    let accept = state.is_routable()
747                        || (include_down && matches!(state, crate::cluster::peer::PeerState::Down));
748                    if accept {
749                        routable.push((dc_idx, rack_idx, peer_idx));
750                    }
751                }
752            }
753        }
754    }
755    routable
756}
757
758fn build_target(
759    dcs: &[crate::cluster::Datacenter],
760    peers: &[crate::cluster::peer::Peer],
761    dc_idx: usize,
762    rack_idx: usize,
763    peer_idx: u32,
764) -> ReplicaTarget {
765    let dc_name = dcs[dc_idx].name().to_string();
766    let rack_name = dcs[dc_idx].racks()[rack_idx].name().to_string();
767    let is_local = peers
768        .get(peer_idx as usize)
769        .is_some_and(crate::cluster::peer::Peer::is_local);
770    ReplicaTarget {
771        peer_idx,
772        dc: dc_name,
773        rack: rack_name,
774        is_local,
775    }
776}
777
778/// The routable replica set for a key, partitioned into local-DC and
779/// remote-DC candidates by [`ClusterDispatcher::plan`].
780struct RoutablePartition {
781    local: Vec<(usize, usize, u32)>,
782    remote: Vec<(usize, usize, u32)>,
783}
784
785fn plan_with_consistency(
786    cfg: &crate::cluster::pool::PoolConfig,
787    dcs: &[crate::cluster::Datacenter],
788    peers: &[crate::cluster::peer::Peer],
789    consistency: ConsistencyLevel,
790    routing: MsgRouting,
791    is_read: bool,
792    partition: RoutablePartition,
793) -> DispatchPlan {
794    let RoutablePartition { local, remote } = partition;
795    let want_per_dc_fanout = matches!(consistency, ConsistencyLevel::DcEachSafeQuorum)
796        || matches!(routing, MsgRouting::AllNodesAllRacksAllDcs);
797    let mut targets: Vec<ReplicaTarget> = Vec::new();
798    match consistency {
799        ConsistencyLevel::DcOne => {
800            if local.is_empty() {
801                return DispatchPlan::NoTargets;
802            }
803            if is_read {
804                // DC_ONE READ: pick the single rack-closest local-DC
805                // replica (lowest latency); the coalescer returns on
806                // its first reply.
807                let mut best: Option<(RackDistance, (usize, usize, u32))> = None;
808                for (dc_idx, rack_idx, peer_idx) in local {
809                    let rack_name = dcs[dc_idx].racks()[rack_idx].name();
810                    let d = rack_distance(&cfg.dc, &cfg.rack, &cfg.dc, rack_name);
811                    let take = match best {
812                        None => true,
813                        Some((bd, _)) => d.cost() < bd.cost(),
814                    };
815                    if take {
816                        best = Some((d, (dc_idx, rack_idx, peer_idx)));
817                    }
818                }
819                if let Some((_, (dc_idx, rack_idx, peer_idx))) = best {
820                    let is_local_node = peers
821                        .get(peer_idx as usize)
822                        .is_some_and(crate::cluster::peer::Peer::is_local);
823                    if is_local_node {
824                        return DispatchPlan::LocalDatastore;
825                    }
826                    targets.push(build_target(dcs, peers, dc_idx, rack_idx, peer_idx));
827                }
828            } else {
829                // DC_ONE WRITE: fan out to EVERY local-DC replica for
830                // durability (each rack is a replica). The coalescer
831                // acks on the first successful reply, but the write is
832                // delivered to all replicas. This matches C Dynomite,
833                // where a DC_ONE write replicates within the local DC.
834                for (dc_idx, rack_idx, peer_idx) in local {
835                    targets.push(build_target(dcs, peers, dc_idx, rack_idx, peer_idx));
836                }
837                // A lone local replica (single-node pool, or any key
838                // whose only local-DC replica is this node) is served
839                // through the local datastore directly. Fanning it out
840                // as a one-element `Replicas` plan would require a wired
841                // backend channel; the `LocalDatastore` path is the one
842                // both the wire proxy and `inject_request` already
843                // serve via the local-datastore hook. Multi-replica
844                // fan-out (more than one local target) is unaffected.
845                if targets.len() == 1 && targets[0].is_local {
846                    return DispatchPlan::LocalDatastore;
847                }
848            }
849        }
850        ConsistencyLevel::DcQuorum | ConsistencyLevel::DcSafeQuorum => {
851            if local.is_empty() {
852                return DispatchPlan::NoTargets;
853            }
854            for (dc_idx, rack_idx, peer_idx) in local {
855                targets.push(build_target(dcs, peers, dc_idx, rack_idx, peer_idx));
856            }
857        }
858        ConsistencyLevel::DcEachSafeQuorum => {
859            if local.is_empty() && remote.is_empty() {
860                return DispatchPlan::NoTargets;
861            }
862            for (dc_idx, rack_idx, peer_idx) in local.iter().chain(remote.iter()) {
863                targets.push(build_target(dcs, peers, *dc_idx, *rack_idx, *peer_idx));
864            }
865        }
866    }
867    if want_per_dc_fanout && !remote.is_empty() {
868        for (dc_idx, rack_idx, peer_idx) in remote {
869            if !targets.iter().any(|t| t.peer_idx == peer_idx) {
870                targets.push(build_target(dcs, peers, dc_idx, rack_idx, peer_idx));
871            }
872        }
873    }
874    if targets.is_empty() {
875        return DispatchPlan::LocalDatastore;
876    }
877    DispatchPlan::Replicas {
878        targets,
879        consistency,
880    }
881}
882
883impl Dispatcher for ClusterDispatcher {
884    #[allow(
885        clippy::too_many_lines,
886        reason = "single dispatch fn must enumerate every plan; splitting hides the planner-to-effect mapping"
887    )]
888    fn dispatch(&self, req: Msg, responder: ServerSink) -> DispatchOutcome {
889        if req.flags().quit {
890            return DispatchOutcome::Drop;
891        }
892        // FT.* / HSET interception. Runs before the routing
893        // planner so vector-index commands never visit the
894        // backend, and HSETs against an indexed prefix get the
895        // vector mirrored into the in-process registry before
896        // the standard storage write fans out.
897        if let Some(ext) = self.command_extension.as_ref() {
898            if let Some(outcome) = self.intercept_command(ext.as_ref(), &req) {
899                return outcome;
900            }
901        }
902        // Inspect the request without consuming it: pull the routing
903        // bytes from the first parsed key. `KeyPos::tag_bytes` returns
904        // the hash-tag-aware sub-range when one was parsed and the full
905        // key otherwise, which is the slice shape `plan` expects.
906        // Requests with no parsed keys (e.g. PING, INFO) fall through
907        // with an empty slice; `plan` handles that by routing to the
908        // local datastore.
909        let key: Vec<u8> = req
910            .keys()
911            .first()
912            .map(|kp| kp.tag_bytes().to_vec())
913            .unwrap_or_default();
914        let plan = self.plan(&req, &key);
915        let (req_span, _plan_span) = enter_plan_span(req.id(), &plan);
916        match plan {
917            DispatchPlan::Drop => DispatchOutcome::Drop,
918            DispatchPlan::NoTargets => {
919                let err_type = if matches!(req.ty(), MsgType::ReqRedisGet | MsgType::ReqRedisSet) {
920                    MsgType::RspRedisError
921                } else {
922                    MsgType::RspMcServerError
923                };
924                let rsp = crate::msg::response::make_error(
925                    &req,
926                    err_type,
927                    0,
928                    crate::msg::DynErrorCode::DynomiteNoQuorumAchieved,
929                    &self.mbuf_pool,
930                );
931                DispatchOutcome::Error(rsp)
932            }
933            DispatchPlan::LocalDatastore => {
934                // In-process datastore hook takes precedence: hand
935                // the parsed request to the embedder's `Datastore`
936                // and deliver its reply on the connection's
937                // responder. The hook is async and `dispatch` is
938                // sync, so spawn the await and return `Pending`
939                // (the same contract the backend-channel and
940                // replica paths use).
941                if let Some(ds) = self.local_datastore.as_ref() {
942                    let ds = Arc::clone(ds);
943                    let req_id = req.id();
944                    let span = req_span.clone();
945                    tokio::spawn(async move {
946                        if let Ok(rsp) = ds.dispatch(req).await {
947                            let _ = responder
948                                .send(OutboundEnvelope {
949                                    req_id,
950                                    rsp,
951                                    span,
952                                    source_peer_idx: None,
953                                })
954                                .await;
955                        }
956                    });
957                    return DispatchOutcome::Pending;
958                }
959                if let Some(tx) = self.backend.as_ref() {
960                    // Snapshot the wire bytes from the parsed mbuf
961                    // chain. The chain is the original on-the-wire
962                    // sequence the parser walked, so this is a
963                    // faithful relay rather than a re-encode.
964                    let bytes: Vec<u8> = req
965                        .mbufs()
966                        .iter()
967                        .flat_map(|b| b.readable().to_vec())
968                        .collect();
969                    if bytes.is_empty() {
970                        // Parsed request with no replayable bytes
971                        // (e.g. a synthetic `Msg`) - drop rather
972                        // than enqueue a no-op on the backend.
973                        return DispatchOutcome::Drop;
974                    }
975                    let env = OutboundRequest {
976                        bytes,
977                        req_id: req.id(),
978                        responder,
979                        span: req_span.clone(),
980                        ty: crate::proto::dnode::DmsgType::Req,
981                        target_peer_idx: None,
982                    };
983                    if let Err(err) = tx.try_send(env) {
984                        // Backend channel full or closed: surface
985                        // an error to the client immediately.
986                        if let Some(m) = self.failure_metrics.as_ref() {
987                            match err {
988                                tokio::sync::mpsc::error::TrySendError::Full(_) => {
989                                    m.record_backend_send_full();
990                                }
991                                tokio::sync::mpsc::error::TrySendError::Closed(_) => {
992                                    m.record_backend_send_closed();
993                                }
994                            }
995                        }
996                        let err_type =
997                            if matches!(req.ty(), MsgType::ReqRedisGet | MsgType::ReqRedisSet) {
998                                MsgType::RspRedisError
999                            } else {
1000                                MsgType::RspMcServerError
1001                            };
1002                        let rsp = crate::msg::response::make_error(
1003                            &req,
1004                            err_type,
1005                            0,
1006                            crate::msg::DynErrorCode::DynomiteNoQuorumAchieved,
1007                            &self.mbuf_pool,
1008                        );
1009                        return DispatchOutcome::Error(rsp);
1010                    }
1011                }
1012                DispatchOutcome::Pending
1013            }
1014            DispatchPlan::Replicas {
1015                targets,
1016                consistency,
1017            } => self.dispatch_replicas(&req, &req_span, &targets, consistency, responder),
1018        }
1019    }
1020}
1021
1022impl ClusterDispatcher {
1023    /// Fan a request out across replicas and install the per-
1024    /// request reply coalescer.
1025    ///
1026    /// The single-target case short-circuits to a direct
1027    /// forward (no coalescer needed). The multi-target case
1028    /// spawns a coalescer task on the ambient tokio runtime; the
1029    /// task drains the per-target replies, picks one according
1030    /// to the consistency level, and forwards the chosen reply
1031    /// to the original `responder`. Divergent replicas are
1032    /// scheduled for read-repair writes via the same
1033    /// `peer_backends` channels.
1034    ///
1035    /// When [`hinted_handoff_active`](Self::hinted_handoff_active)
1036    /// reports true and the request is a write, targets that
1037    /// are currently in [`crate::cluster::peer::PeerState::Down`]
1038    /// are recorded in the hint store instead of being sent.
1039    /// A synthetic `+OK\r\n` reply is fed to the coalescer on
1040    /// the hinted target's behalf so the consistency threshold
1041    /// can be met by the surviving replicas plus the hint(s).
1042    /// On `try_send` failure (channel closed or full) for a
1043    /// non-Down peer, the dispatcher likewise falls back to
1044    /// hinting before declaring the target lost.
1045    fn dispatch_replicas(
1046        &self,
1047        req: &Msg,
1048        req_span: &tracing::Span,
1049        targets: &[ReplicaTarget],
1050        consistency: ConsistencyLevel,
1051        responder: ServerSink,
1052    ) -> DispatchOutcome {
1053        if targets.is_empty() {
1054            return DispatchOutcome::Drop;
1055        }
1056        // Snapshot the wire bytes once. Each target gets its
1057        // own clone (the ServerConn / DnodeServerConn takes
1058        // ownership of `bytes`).
1059        let bytes: Vec<u8> = req
1060            .mbufs()
1061            .iter()
1062            .flat_map(|b| b.readable().to_vec())
1063            .collect();
1064        if bytes.is_empty() {
1065            return DispatchOutcome::Drop;
1066        }
1067        // Snapshot the per-target current state so we do not
1068        // re-acquire the peer-table lock inside the per-target
1069        // loop. Local peers are always treated as Normal.
1070        let peer_states = self.snapshot_peer_states(targets);
1071        let is_read = matches!(req.ty(), MsgType::Unknown) || req.flags().is_read;
1072        let is_write = !is_read;
1073        let handoff_active = self.hinted_handoff_active() && is_write;
1074        // Single-target path: no coalescing needed; forward
1075        // directly to the original responder.
1076        if targets.len() == 1 {
1077            return self.dispatch_replicas_direct(
1078                req,
1079                req_span,
1080                targets,
1081                &bytes,
1082                &responder,
1083                &HandoffCtx {
1084                    handoff_active,
1085                    peer_states: &peer_states,
1086                },
1087            );
1088        }
1089        // Multi-target path: install the coalescer.
1090        let cfg = self.pool.config();
1091        let local_dc = cfg.dc.clone();
1092        // Channel each replica's reply lands on. Sized to
1093        // `targets.len() + 1`: every target produces at most one
1094        // envelope (real or hint-synthesised) and we leave a
1095        // spare so a late repair reply never blocks the actor.
1096        let (intermediate_tx, intermediate_rx) =
1097            mpsc::channel::<OutboundEnvelope>(targets.len() + 1);
1098        // Build the tracker's target list and capture the
1099        // per-target dispatch state.
1100        let target_pairs: Vec<(u32, String)> =
1101            targets.iter().map(|t| (t.peer_idx, t.dc.clone())).collect();
1102        // Read repair context: the original primary key (single-
1103        // key requests only) and the request type.
1104        let repair_key: Option<Vec<u8>> = req
1105            .keys()
1106            .first()
1107            .map(|kp| kp.tag_bytes().to_vec())
1108            .filter(|k| !k.is_empty());
1109        let repair_ctx = repair_key.map(|key| ReadRepairContext {
1110            req_id: req.id(),
1111            req_ty: req.ty(),
1112            key,
1113            mbuf_pool: self.mbuf_pool.clone(),
1114            peer_backends: self.peer_backends.clone(),
1115            local_backend: self.backend.clone(),
1116            target_is_local: targets.iter().map(|t| (t.peer_idx, t.is_local)).collect(),
1117        });
1118        // Fan out: each per-target outbound feeds the coalescer
1119        // channel, NOT the client's responder.
1120        let mut sent = 0usize;
1121        let mut hinted = 0usize;
1122        for target in targets {
1123            let action = Self::choose_target_action(target, handoff_active, &peer_states);
1124            match action {
1125                TargetAction::Send => {
1126                    if self.fanout_send(target, req, req_span, &bytes, &intermediate_tx) {
1127                        sent += 1;
1128                    } else if handoff_active
1129                        && self.hint_target(target, &bytes, req, req_span, &intermediate_tx)
1130                    {
1131                        hinted += 1;
1132                    }
1133                }
1134                TargetAction::Hint => {
1135                    if self.hint_target(target, &bytes, req, req_span, &intermediate_tx) {
1136                        hinted += 1;
1137                    }
1138                }
1139            }
1140        }
1141        // Drop the local clone of the intermediate sender so the
1142        // coalescer task observes RX close once every per-target
1143        // sender has been dropped. (`OutboundRequest` owns one
1144        // sender each; once they finish they drop it.)
1145        drop(intermediate_tx);
1146        if sent + hinted == 0 {
1147            return DispatchOutcome::Error(self.no_quorum_error(req));
1148        }
1149        let req_id = req.id();
1150        let req_ty = req.ty();
1151        let mbuf_pool = self.mbuf_pool.clone();
1152        let failure_metrics = self.failure_metrics.clone();
1153        tokio::spawn(coalesce_actor(
1154            req_id,
1155            req_ty,
1156            consistency,
1157            target_pairs,
1158            local_dc,
1159            intermediate_rx,
1160            responder,
1161            mbuf_pool,
1162            repair_ctx,
1163            failure_metrics,
1164        ));
1165        DispatchOutcome::Pending
1166    }
1167
1168    /// Capture the current `PeerState` for each target so the
1169    /// per-target loop does not re-acquire the read lock on
1170    /// every iteration. Local-node targets are reported as
1171    /// `Normal` regardless of the on-disk peer entry.
1172    fn snapshot_peer_states(
1173        &self,
1174        targets: &[ReplicaTarget],
1175    ) -> std::collections::HashMap<u32, crate::cluster::peer::PeerState> {
1176        use crate::cluster::peer::PeerState;
1177        let peers = self.pool.peers().read();
1178        let mut out = std::collections::HashMap::with_capacity(targets.len());
1179        for t in targets {
1180            let state = if t.is_local {
1181                PeerState::Normal
1182            } else {
1183                peers
1184                    .get(t.peer_idx as usize)
1185                    .map_or(PeerState::Unknown, crate::cluster::peer::Peer::state)
1186            };
1187            out.insert(t.peer_idx, state);
1188        }
1189        out
1190    }
1191
1192    fn choose_target_action(
1193        target: &ReplicaTarget,
1194        handoff_active: bool,
1195        peer_states: &std::collections::HashMap<u32, crate::cluster::peer::PeerState>,
1196    ) -> TargetAction {
1197        use crate::cluster::peer::PeerState;
1198        if !handoff_active {
1199            return TargetAction::Send;
1200        }
1201        let state = peer_states
1202            .get(&target.peer_idx)
1203            .copied()
1204            .unwrap_or(PeerState::Unknown);
1205        match state {
1206            PeerState::Down => TargetAction::Hint,
1207            _ => TargetAction::Send,
1208        }
1209    }
1210
1211    /// Forward one target via its outbound channel. Returns
1212    /// `true` on a successful `try_send`.
1213    fn fanout_send(
1214        &self,
1215        target: &ReplicaTarget,
1216        req: &Msg,
1217        req_span: &tracing::Span,
1218        bytes: &[u8],
1219        intermediate_tx: &mpsc::Sender<OutboundEnvelope>,
1220    ) -> bool {
1221        // A request forwarded to a REMOTE peer is tagged
1222        // `ReqForward` so the receiver hands it straight to its
1223        // local datastore instead of re-routing it (which would
1224        // re-hash, re-plan, and in the worst case bounce the
1225        // request back). The local-backend path keeps `Req`.
1226        let ty = if target.is_local {
1227            crate::proto::dnode::DmsgType::Req
1228        } else {
1229            crate::proto::dnode::DmsgType::ReqForward
1230        };
1231        let env = OutboundRequest {
1232            bytes: bytes.to_vec(),
1233            req_id: req.id(),
1234            responder: intermediate_tx.clone(),
1235            span: req_span.clone(),
1236            ty,
1237            target_peer_idx: Some(target.peer_idx),
1238        };
1239        let send_result = if target.is_local {
1240            self.backend.as_ref().map(|tx| tx.try_send(env))
1241        } else {
1242            self.peer_backends
1243                .get(&target.peer_idx)
1244                .map(|tx| tx.try_send(env))
1245        };
1246        match send_result {
1247            Some(Ok(())) => true,
1248            Some(Err(err)) => {
1249                self.observe_send_error(target, &err);
1250                false
1251            }
1252            None => false,
1253        }
1254    }
1255
1256    /// Convert a `tokio::sync::mpsc::error::TrySendError` into a
1257    /// failure-metrics observation. Local targets bump the
1258    /// backend counters; peer targets bump the per-peer
1259    /// counters labelled with the peer's DC.
1260    fn observe_send_error(
1261        &self,
1262        target: &ReplicaTarget,
1263        err: &tokio::sync::mpsc::error::TrySendError<OutboundRequest>,
1264    ) {
1265        let Some(m) = self.failure_metrics.as_ref() else {
1266            return;
1267        };
1268        if target.is_local {
1269            match err {
1270                tokio::sync::mpsc::error::TrySendError::Full(_) => m.record_backend_send_full(),
1271                tokio::sync::mpsc::error::TrySendError::Closed(_) => {
1272                    m.record_backend_send_closed();
1273                }
1274            }
1275        } else {
1276            let peer_dc = self.peer_dc_label(target.peer_idx);
1277            match err {
1278                tokio::sync::mpsc::error::TrySendError::Full(_) => {
1279                    m.record_peer_send_full(target.peer_idx, &peer_dc);
1280                }
1281                tokio::sync::mpsc::error::TrySendError::Closed(_) => {
1282                    m.record_peer_send_closed(target.peer_idx, &peer_dc);
1283                }
1284            }
1285        }
1286    }
1287
1288    /// Record `bytes` as a hint for `target`'s peer and feed the
1289    /// coalescer a synthetic success reply. Returns `true` when
1290    /// both the enqueue and the synth-push succeeded.
1291    fn hint_target(
1292        &self,
1293        target: &ReplicaTarget,
1294        bytes: &[u8],
1295        req: &Msg,
1296        req_span: &tracing::Span,
1297        intermediate_tx: &mpsc::Sender<OutboundEnvelope>,
1298    ) -> bool {
1299        let Some(store) = self.hint_store.as_ref() else {
1300            return false;
1301        };
1302        let cfg = self.pool.config();
1303        let ttl = std::time::Duration::from_secs(cfg.hint_ttl_seconds.max(1));
1304        match store.enqueue(target.peer_idx, bytes.to_vec(), ttl) {
1305            Ok(()) => {}
1306            Err(e) => {
1307                tracing::debug!(
1308                    target: "dynomite::hints",
1309                    peer_idx = target.peer_idx,
1310                    error = %e,
1311                    "hint enqueue failed"
1312                );
1313                return false;
1314            }
1315        }
1316        let synth = synth_hint_reply(req, &self.mbuf_pool);
1317        let env = OutboundEnvelope {
1318            req_id: req.id(),
1319            rsp: synth,
1320            span: req_span.clone(),
1321            source_peer_idx: Some(target.peer_idx),
1322        };
1323        if intermediate_tx.try_send(env).is_err() {
1324            // The channel is sized for one envelope per target;
1325            // an immediate full means the coalescer task has
1326            // exited. Still report success so the hint sits in
1327            // the store for the drainer.
1328            tracing::debug!(
1329                target: "dynomite::hints",
1330                peer_idx = target.peer_idx,
1331                "hint synth-reply could not be queued; coalescer absent"
1332            );
1333        }
1334        tracing::debug!(
1335            target: "dynomite::hints",
1336            peer_idx = target.peer_idx,
1337            bytes = bytes.len(),
1338            "stored hint for down peer"
1339        );
1340        true
1341    }
1342
1343    fn dispatch_replicas_direct(
1344        &self,
1345        req: &Msg,
1346        req_span: &tracing::Span,
1347        targets: &[ReplicaTarget],
1348        bytes: &[u8],
1349        responder: &ServerSink,
1350        ctx: &HandoffCtx<'_>,
1351    ) -> DispatchOutcome {
1352        debug_assert_eq!(targets.len(), 1);
1353        let target = &targets[0];
1354        // If the target is Down and handoff is active, hint it
1355        // and feed the responder a synth `+OK\r\n` so the
1356        // client sees the write as having succeeded.
1357        if let TargetAction::Hint =
1358            Self::choose_target_action(target, ctx.handoff_active, ctx.peer_states)
1359        {
1360            if self.hint_single_target_direct(target, bytes, req, req_span, responder) {
1361                return DispatchOutcome::Pending;
1362            }
1363            return DispatchOutcome::Error(self.no_quorum_error(req));
1364        }
1365        let env = OutboundRequest {
1366            bytes: bytes.to_vec(),
1367            req_id: req.id(),
1368            responder: responder.clone(),
1369            span: req_span.clone(),
1370            // A forward to a remote peer is `ReqForward` so the
1371            // receiver serves it locally rather than re-routing it;
1372            // the local-backend path uses `Req`.
1373            ty: if target.is_local {
1374                crate::proto::dnode::DmsgType::Req
1375            } else {
1376                crate::proto::dnode::DmsgType::ReqForward
1377            },
1378            target_peer_idx: Some(target.peer_idx),
1379        };
1380        let send_result = if target.is_local {
1381            self.backend.as_ref().map(|tx| tx.try_send(env))
1382        } else {
1383            self.peer_backends
1384                .get(&target.peer_idx)
1385                .map(|tx| tx.try_send(env))
1386        };
1387        let sent = match send_result {
1388            Some(Ok(())) => true,
1389            Some(Err(ref err)) => {
1390                self.observe_send_error(target, err);
1391                false
1392            }
1393            None => false,
1394        };
1395        if sent {
1396            return DispatchOutcome::Pending;
1397        }
1398        if ctx.handoff_active
1399            && self.hint_single_target_direct(target, bytes, req, req_span, responder)
1400        {
1401            return DispatchOutcome::Pending;
1402        }
1403        DispatchOutcome::Error(self.no_quorum_error(req))
1404    }
1405
1406    /// Single-target hint path. Records the hint in the store
1407    /// and pushes a `+OK\r\n` envelope onto the responder. The
1408    /// client sees the write as having succeeded; the drainer
1409    /// will replay the request to the peer when it returns.
1410    fn hint_single_target_direct(
1411        &self,
1412        target: &ReplicaTarget,
1413        bytes: &[u8],
1414        req: &Msg,
1415        req_span: &tracing::Span,
1416        responder: &ServerSink,
1417    ) -> bool {
1418        let Some(store) = self.hint_store.as_ref() else {
1419            return false;
1420        };
1421        let cfg = self.pool.config();
1422        let ttl = std::time::Duration::from_secs(cfg.hint_ttl_seconds.max(1));
1423        if let Err(e) = store.enqueue(target.peer_idx, bytes.to_vec(), ttl) {
1424            tracing::debug!(
1425                target: "dynomite::hints",
1426                peer_idx = target.peer_idx,
1427                error = %e,
1428                "hint enqueue failed (single-target)"
1429            );
1430            return false;
1431        }
1432        let synth = synth_hint_reply(req, &self.mbuf_pool);
1433        let env = OutboundEnvelope {
1434            req_id: req.id(),
1435            rsp: synth,
1436            span: req_span.clone(),
1437            source_peer_idx: Some(target.peer_idx),
1438        };
1439        let _ = responder.try_send(env);
1440        true
1441    }
1442
1443    fn no_quorum_error(&self, req: &Msg) -> Msg {
1444        let err_type = if matches!(req.ty(), MsgType::ReqRedisGet | MsgType::ReqRedisSet) {
1445            MsgType::RspRedisError
1446        } else {
1447            MsgType::RspMcServerError
1448        };
1449        crate::msg::response::make_error(
1450            req,
1451            err_type,
1452            0,
1453            crate::msg::DynErrorCode::DynomiteNoQuorumAchieved,
1454            &self.mbuf_pool,
1455        )
1456    }
1457
1458    /// FT.* / HSET interception. Returns `Some(outcome)` when the
1459    /// command was fully handled (FT.* keyword, or an HSET that
1460    /// references an indexed prefix with a malformed payload);
1461    /// returns `None` to let the caller fall through to the
1462    /// standard dispatch path. The HSET success path returns
1463    /// `None` after the extension records the side-effect so
1464    /// the standard storage write still goes to the backend.
1465    fn intercept_command(
1466        &self,
1467        ext: &dyn crate::embed::CommandExtension,
1468        req: &Msg,
1469    ) -> Option<DispatchOutcome> {
1470        if ext.handles_msg_type(req.ty()) {
1471            return Some(self.run_extension_command(ext, req));
1472        }
1473        if matches!(req.ty(), MsgType::ReqRedisHset) {
1474            return self.intercept_extension_hset(ext, req);
1475        }
1476        None
1477    }
1478
1479    /// Drive an FT.* command through the registered extension
1480    /// and wrap the RESP bytes in a [`DispatchOutcome::Inline`].
1481    /// When the extension declines (`try_dispatch` returns
1482    /// `None`) the outcome is rendered as a `-ERR not
1483    /// supported in this build` reply so the dispatcher does
1484    /// not silently drop the request.
1485    fn run_extension_command(
1486        &self,
1487        ext: &dyn crate::embed::CommandExtension,
1488        req: &Msg,
1489    ) -> DispatchOutcome {
1490        // For typed FT.* variants the keyword is unambiguous; for
1491        // [`MsgType::ReqRedisFtUnknown`] we recover the original
1492        // wire keyword from the request's mbuf chain (the parser
1493        // case-folds for table lookup but the wire bytes are
1494        // preserved verbatim) so the extension can decide
1495        // whether we recognise the keyword and either execute
1496        // it or surface a structured `-ERR ...` reply.
1497        let recovered_kw: Vec<u8>;
1498        let keyword: &[u8] = match req.ty() {
1499            MsgType::ReqRedisFtCreate => b"FT.CREATE",
1500            MsgType::ReqRedisFtSearch => b"FT.SEARCH",
1501            MsgType::ReqRedisFtInfo => b"FT.INFO",
1502            MsgType::ReqRedisFtList => b"FT.LIST",
1503            MsgType::ReqRedisFtDropindex => b"FT.DROPINDEX",
1504            MsgType::ReqRedisFtRegex => b"FT.REGEX",
1505            MsgType::ReqRedisFtSugadd => b"FT.SUGADD",
1506            MsgType::ReqRedisFtSugget => b"FT.SUGGET",
1507            MsgType::ReqRedisFtSugdel => b"FT.SUGDEL",
1508            MsgType::ReqRedisFtSuglen => b"FT.SUGLEN",
1509            MsgType::ReqRedisFtUnknown => {
1510                recovered_kw = first_bulk_token(req).unwrap_or_else(|| b"FT.UNKNOWN".to_vec());
1511                recovered_kw.as_slice()
1512            }
1513            // The dispatcher only enters this branch from the
1514            // FT.* arm in [`Self::intercept_command`], so the
1515            // catch-all is unreachable unless the MsgType set
1516            // drifts out of sync with that arm.
1517            _ => return DispatchOutcome::Drop,
1518        };
1519        let mut args: Vec<&[u8]> = Vec::with_capacity(1 + req.keys().len() + req.args().len());
1520        args.push(keyword);
1521        for k in req.keys() {
1522            args.push(k.key());
1523        }
1524        for a in req.args() {
1525            args.push(a.bytes());
1526        }
1527        let bytes = ext.try_dispatch(&args).unwrap_or_else(|| {
1528            let kw = String::from_utf8_lossy(keyword);
1529            format!("-ERR not supported in this build: {kw}\r\n").into_bytes()
1530        });
1531        DispatchOutcome::Inline(synthetic_redis_reply(req, &self.mbuf_pool, &bytes))
1532    }
1533
1534    /// HSET interception via the registered extension.
1535    /// Returns `Some(Error(...))` when the extension reports a
1536    /// malformed payload against a registered prefix; returns
1537    /// `None` (so the dispatcher falls through to the backend)
1538    /// on success or when no registered prefix matches.
1539    fn intercept_extension_hset(
1540        &self,
1541        ext: &dyn crate::embed::CommandExtension,
1542        req: &Msg,
1543    ) -> Option<DispatchOutcome> {
1544        let mut args: Vec<&[u8]> = Vec::with_capacity(req.keys().len() + req.args().len());
1545        for k in req.keys() {
1546            args.push(k.key());
1547        }
1548        for a in req.args() {
1549            args.push(a.bytes());
1550        }
1551        match ext.try_intercept_hset(&args) {
1552            crate::embed::HsetOutcome::Absorbed | crate::embed::HsetOutcome::NotIndexed => None,
1553            crate::embed::HsetOutcome::Error(message) => {
1554                let payload = format!("-ERR {message}\r\n");
1555                Some(DispatchOutcome::Error(synthetic_redis_reply(
1556                    req,
1557                    &self.mbuf_pool,
1558                    payload.as_bytes(),
1559                )))
1560            }
1561        }
1562    }
1563}
1564
1565/// Wrap an arbitrary RESP byte sequence as a synthetic Redis
1566/// response [`Msg`]. The response inherits the request's id (so
1567/// the FSM can pair it with the originating request), is marked
1568/// `is_request = false`, and the supplied bytes are copied into
1569/// one or more mbufs drawn from `pool`.
1570fn synthetic_redis_reply(req: &Msg, pool: &MbufPool, payload: &[u8]) -> Msg {
1571    let mut rsp = Msg::new(req.id(), MsgType::RspRedisStatus, false);
1572    rsp.set_parent_id(req.id());
1573    let mut written = 0usize;
1574    while written < payload.len() {
1575        let mut buf = pool.get();
1576        let n = buf.recv(&payload[written..]);
1577        debug_assert!(
1578            n > 0,
1579            "MbufPool returned a buffer with zero writable capacity"
1580        );
1581        rsp.mbufs_mut().push_back(buf);
1582        written += n;
1583    }
1584    rsp.recompute_mlen();
1585    rsp
1586}
1587
1588/// Recover the first bulk-string token from a parsed RESP
1589/// request's mbuf chain. The parser case-folds the keyword for
1590/// the command-table lookup but stores the original wire bytes
1591/// in the mbufs; this helper reads them back so the dispatcher
1592/// can render structured error replies that quote the actual
1593/// keyword the client sent.
1594///
1595/// Returns `None` when the mbuf chain does not begin with a
1596/// well-formed `*N\r\n$M\r\n<token>\r\n` sequence.
1597fn first_bulk_token(req: &Msg) -> Option<Vec<u8>> {
1598    let mut wire: Vec<u8> = Vec::new();
1599    for buf in req.mbufs() {
1600        wire.extend_from_slice(buf.readable());
1601        if wire.len() > 256 {
1602            break;
1603        }
1604    }
1605    let mut p = 0usize;
1606    if wire.first() == Some(&b'*') {
1607        let cr = wire.iter().position(|&b| b == b'\r')?;
1608        if wire.get(cr + 1) != Some(&b'\n') {
1609            return None;
1610        }
1611        p = cr + 2;
1612    }
1613    if wire.get(p) != Some(&b'$') {
1614        return None;
1615    }
1616    let header_start = p + 1;
1617    let header_cr = wire[header_start..]
1618        .iter()
1619        .position(|&b| b == b'\r')
1620        .map(|i| header_start + i)?;
1621    if wire.get(header_cr + 1) != Some(&b'\n') {
1622        return None;
1623    }
1624    let len_str = std::str::from_utf8(&wire[header_start..header_cr]).ok()?;
1625    let len: usize = len_str.parse().ok()?;
1626    let body_start = header_cr + 2;
1627    let body_end = body_start.checked_add(len)?;
1628    if wire.len() < body_end + 2 {
1629        return None;
1630    }
1631    Some(wire[body_start..body_end].to_vec())
1632}
1633
1634/// Context required to schedule read-repair writes once the
1635/// coalescer has identified a winner and a divergent set.
1636#[derive(Clone)]
1637struct ReadRepairContext {
1638    req_id: crate::core::types::MsgId,
1639    req_ty: MsgType,
1640    /// Original primary key (single-key requests only). The v1
1641    /// repair scheduler operates over single-key Redis reads;
1642    /// multi-key fragmentation goes through a separate path.
1643    key: Vec<u8>,
1644    mbuf_pool: MbufPool,
1645    peer_backends: std::collections::HashMap<u32, mpsc::Sender<OutboundRequest>>,
1646    local_backend: Option<mpsc::Sender<OutboundRequest>>,
1647    target_is_local: std::collections::HashMap<u32, bool>,
1648}
1649
1650/// Per-fan-out coalescer task body.
1651#[allow(
1652    clippy::too_many_arguments,
1653    reason = "actor task captures the entire dispatch context; bundling into a struct adds churn for no callsite gain"
1654)]
1655async fn coalesce_actor(
1656    req_id: crate::core::types::MsgId,
1657    req_ty: MsgType,
1658    consistency: ConsistencyLevel,
1659    targets: Vec<(u32, String)>,
1660    local_dc: String,
1661    mut intermediate_rx: mpsc::Receiver<OutboundEnvelope>,
1662    client_tx: ServerSink,
1663    mbuf_pool: MbufPool,
1664    repair_ctx: Option<ReadRepairContext>,
1665    failure_metrics: Option<Arc<crate::stats::FailureMetrics>>,
1666) {
1667    use crate::proto::redis::{CoalesceOutcome, CoalesceTracker};
1668    let mut tracker = CoalesceTracker::new(req_id, consistency, targets, &local_dc);
1669    let mut emitted = false;
1670    while let Some(env) = intermediate_rx.recv().await {
1671        let source = env.source_peer_idx.unwrap_or(u32::MAX);
1672        let span = env.span.clone();
1673        let outcome = tracker.record_reply(source, env.rsp);
1674        match outcome {
1675            CoalesceOutcome::Pending => {}
1676            CoalesceOutcome::Ready {
1677                winner,
1678                divergent_targets,
1679            } => {
1680                if !emitted {
1681                    let winner_bytes: Vec<u8> = winner
1682                        .mbufs()
1683                        .iter()
1684                        .flat_map(|b| b.readable().to_vec())
1685                        .collect();
1686                    let out_env = OutboundEnvelope {
1687                        req_id,
1688                        rsp: *winner,
1689                        span: span.clone(),
1690                        source_peer_idx: None,
1691                    };
1692                    let _ = client_tx.send(out_env).await;
1693                    emitted = true;
1694                    if !divergent_targets.is_empty() {
1695                        if let Some(ctx) = repair_ctx.as_ref() {
1696                            schedule_read_repair(ctx, &divergent_targets, &winner_bytes, &span);
1697                        }
1698                    }
1699                }
1700            }
1701            CoalesceOutcome::Error(reason) => {
1702                if !emitted {
1703                    let err_type = if matches!(req_ty, MsgType::ReqRedisGet | MsgType::ReqRedisSet)
1704                    {
1705                        MsgType::RspRedisError
1706                    } else {
1707                        MsgType::RspMcServerError
1708                    };
1709                    let anchor = Msg::new(req_id, req_ty, true);
1710                    let rsp = crate::msg::response::make_error(
1711                        &anchor,
1712                        err_type,
1713                        0,
1714                        crate::msg::DynErrorCode::DynomiteNoQuorumAchieved,
1715                        &mbuf_pool,
1716                    );
1717                    let _ = client_tx
1718                        .send(OutboundEnvelope {
1719                            req_id,
1720                            rsp,
1721                            span: span.clone(),
1722                            source_peer_idx: None,
1723                        })
1724                        .await;
1725                    emitted = true;
1726                }
1727                tracing::debug!(target: "dynomite::coalesce", req_id, reason = %reason, "coalesce error");
1728            }
1729        }
1730    }
1731    if !emitted {
1732        // No reply was emitted and the channel closed (every
1733        // per-target sender dropped without producing a reply).
1734        // From the dispatcher's perspective the request has
1735        // timed out at the coalescer layer; surface a
1736        // quorum-unreachable error so the client does not hang
1737        // and bump the response-timeout counter.
1738        if let Some(m) = failure_metrics.as_ref() {
1739            m.record_response_timeout(consistency);
1740        }
1741        let err_type = if matches!(req_ty, MsgType::ReqRedisGet | MsgType::ReqRedisSet) {
1742            MsgType::RspRedisError
1743        } else {
1744            MsgType::RspMcServerError
1745        };
1746        let anchor = Msg::new(req_id, req_ty, true);
1747        let rsp = crate::msg::response::make_error(
1748            &anchor,
1749            err_type,
1750            0,
1751            crate::msg::DynErrorCode::DynomiteNoQuorumAchieved,
1752            &mbuf_pool,
1753        );
1754        let _ = client_tx
1755            .send(OutboundEnvelope {
1756                req_id,
1757                rsp,
1758                span: tracing::Span::none(),
1759                source_peer_idx: None,
1760            })
1761            .await;
1762    }
1763}
1764
1765/// Build a sink the read-repair task can drop replies into. The
1766/// scheduler is fire-and-forget: every reply is discarded.
1767fn repair_sink() -> ServerSink {
1768    let (tx, mut rx) = mpsc::channel::<OutboundEnvelope>(8);
1769    tokio::spawn(async move {
1770        while rx.recv().await.is_some() {
1771            // Drop the envelope; the original client already
1772            // received its reply on the main responder.
1773        }
1774    });
1775    tx
1776}
1777
1778/// Decode a winning RESP reply into the bytes we want to write
1779/// back to divergent replicas.
1780///
1781/// Returns `Some(bytes)` for a bulk-string winner (we ship a
1782/// `SET key value` to the divergent replica) or for a nil reply
1783/// (we ship a `DEL key`). Returns `None` for any other shape
1784/// (errors, integers, multibulk, ...) since the v1 repair
1785/// scheduler only handles single-bulk Redis GET-style winners.
1786fn decode_winner_for_repair(payload: &[u8]) -> Option<RepairAction> {
1787    if payload == b"$-1\r\n" {
1788        return Some(RepairAction::Delete);
1789    }
1790    if !payload.starts_with(b"$") {
1791        return None;
1792    }
1793    // `$<len>\r\n<value>\r\n`
1794    let crlf = payload.iter().position(|&b| b == b'\r')?;
1795    if payload.get(crlf + 1).copied() != Some(b'\n') {
1796        return None;
1797    }
1798    let len_str = std::str::from_utf8(&payload[1..crlf]).ok()?;
1799    let len: usize = len_str.parse().ok()?;
1800    let body_start = crlf + 2;
1801    let body_end = body_start.checked_add(len)?;
1802    if payload.len() < body_end + 2 {
1803        return None;
1804    }
1805    if &payload[body_end..body_end + 2] != b"\r\n" {
1806        return None;
1807    }
1808    Some(RepairAction::Write(payload[body_start..body_end].to_vec()))
1809}
1810
1811/// Bundle of handoff state passed into
1812/// [`ClusterDispatcher::dispatch_replicas_direct`]. Avoids a
1813/// noisy parameter list while keeping the per-call decisions
1814/// ("is handoff on?", "what state is each target in?") in one
1815/// place.
1816struct HandoffCtx<'a> {
1817    handoff_active: bool,
1818    peer_states: &'a std::collections::HashMap<u32, crate::cluster::peer::PeerState>,
1819}
1820
1821/// Per-target dispatch action chosen by
1822/// [`ClusterDispatcher::choose_target_action`].
1823#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1824enum TargetAction {
1825    /// Forward the request via the per-peer outbound channel.
1826    Send,
1827    /// Record a hint and feed the coalescer a synth success
1828    /// reply.
1829    Hint,
1830}
1831
1832/// Synthetic reply pushed into the coalescer on behalf of a
1833/// hinted target. Always a `+OK\r\n` Redis status reply: the
1834/// dispatcher decouples the hint from the eventual peer reply
1835/// (the drainer will fire-and-forget the hint when the peer
1836/// returns), so the synth shape only needs to satisfy the
1837/// coalescer for SET-style writes which is the dominant case
1838/// for hinted handoff. Mismatched-shape writes (DEL with one
1839/// hinted target) coalesce to the surviving real reply via the
1840/// plurality / quorum branch and the divergence is harmless
1841/// because read-repair only fires for `MsgType::ReqRedisGet`.
1842fn synth_hint_reply(req: &Msg, pool: &MbufPool) -> Msg {
1843    crate::msg::response::make_simple_redis(req, pool, b"+OK\r\n")
1844}
1845
1846/// Action a read-repair task should take against a divergent
1847/// replica.
1848enum RepairAction {
1849    /// Ship `SET key <bytes>` to overwrite the stale value.
1850    Write(Vec<u8>),
1851    /// Ship `DEL key` to drop the stale value (winning reply
1852    /// was a nil bulk).
1853    Delete,
1854}
1855
1856/// Build the wire bytes for a Redis repair write.
1857fn build_repair_bytes(action: &RepairAction, key: &[u8]) -> Vec<u8> {
1858    match action {
1859        RepairAction::Write(value) => {
1860            let mut out = Vec::with_capacity(key.len() + value.len() + 32);
1861            out.extend_from_slice(b"*3\r\n$3\r\nSET\r\n$");
1862            out.extend_from_slice(key.len().to_string().as_bytes());
1863            out.extend_from_slice(b"\r\n");
1864            out.extend_from_slice(key);
1865            out.extend_from_slice(b"\r\n$");
1866            out.extend_from_slice(value.len().to_string().as_bytes());
1867            out.extend_from_slice(b"\r\n");
1868            out.extend_from_slice(value);
1869            out.extend_from_slice(b"\r\n");
1870            out
1871        }
1872        RepairAction::Delete => {
1873            let mut out = Vec::with_capacity(key.len() + 24);
1874            out.extend_from_slice(b"*2\r\n$3\r\nDEL\r\n$");
1875            out.extend_from_slice(key.len().to_string().as_bytes());
1876            out.extend_from_slice(b"\r\n");
1877            out.extend_from_slice(key);
1878            out.extend_from_slice(b"\r\n");
1879            out
1880        }
1881    }
1882}
1883
1884/// Schedule fire-and-forget read-repair writes to every
1885/// divergent target. The function only awaits a bounded mpsc
1886/// permit; it never blocks for the repair to complete or for
1887/// the divergent replica to ack.
1888///
1889/// The repair shape is decoded from `winner_bytes`:
1890///
1891/// * Bulk-string winner -> `SET key <value>`.
1892/// * Nil bulk winner -> `DEL key`.
1893/// * Anything else -> skipped (entropy reconciliation will
1894///   handle it later). This v1 limitation is documented in the
1895///   dispatcher tests and in `docs/parity.md`.
1896///
1897/// Repair writes are tagged with `DmsgType::ReqForward` so the
1898/// receiving peer's `dnode_client_loop` rewrites the parsed
1899/// request's routing tag to `LocalNodeOnly`, preventing a
1900/// recursive multi-replica fan-out at the divergent peer.
1901fn schedule_read_repair(
1902    ctx: &ReadRepairContext,
1903    divergent: &[u32],
1904    winner_bytes: &[u8],
1905    span: &tracing::Span,
1906) {
1907    if !matches!(ctx.req_ty, MsgType::ReqRedisGet) {
1908        return;
1909    }
1910    let Some(action) = decode_winner_for_repair(winner_bytes) else {
1911        return;
1912    };
1913    let bytes = build_repair_bytes(&action, &ctx.key);
1914    let sink = repair_sink();
1915    for &peer_idx in divergent {
1916        let is_local = ctx.target_is_local.get(&peer_idx).copied().unwrap_or(false);
1917        let env = OutboundRequest {
1918            bytes: bytes.clone(),
1919            req_id: ctx.req_id,
1920            responder: sink.clone(),
1921            span: span.clone(),
1922            ty: crate::proto::dnode::DmsgType::ReqForward,
1923            target_peer_idx: Some(peer_idx),
1924        };
1925        let sent = if is_local {
1926            ctx.local_backend
1927                .as_ref()
1928                .is_some_and(|tx| tx.try_send(env).is_ok())
1929        } else {
1930            ctx.peer_backends
1931                .get(&peer_idx)
1932                .is_some_and(|tx| tx.try_send(env).is_ok())
1933        };
1934        if sent {
1935            let _ = &ctx.mbuf_pool;
1936            tracing::debug!(
1937                target: "dynomite::read_repair",
1938                req_id = ctx.req_id,
1939                peer_idx,
1940                bytes = bytes.len(),
1941                "scheduled read-repair write",
1942            );
1943        } else {
1944            tracing::debug!(
1945                target: "dynomite::read_repair",
1946                req_id = ctx.req_id,
1947                peer_idx,
1948                "read-repair drop: backend channel unavailable or full",
1949            );
1950        }
1951    }
1952}
1953
1954#[cfg(test)]
1955mod tests {
1956    use super::*;
1957    use crate::cluster::peer::{Peer, PeerEndpoint, PeerState};
1958    use crate::conf::DataStore;
1959    use crate::hashkit::DynToken;
1960
1961    fn cfg(read: ConsistencyLevel, write: ConsistencyLevel) -> crate::cluster::PoolConfig {
1962        crate::cluster::PoolConfig {
1963            read_consistency: read,
1964            write_consistency: write,
1965            dc: "dc1".into(),
1966            rack: "rA".into(),
1967            ..crate::cluster::PoolConfig::default()
1968        }
1969    }
1970
1971    fn peer(idx: u32, dc: &str, rack: &str, tok: u32, is_local: bool, is_same: bool) -> Peer {
1972        let mut p = Peer::new(
1973            idx,
1974            PeerEndpoint::tcp("h".into(), 8101 + u16::try_from(idx).unwrap_or(0)),
1975            rack.into(),
1976            dc.into(),
1977            vec![DynToken::from_u32(tok)],
1978            is_local,
1979            is_same,
1980            false,
1981        );
1982        p.set_state(PeerState::Normal, 0);
1983        p
1984    }
1985
1986    fn pool(read: ConsistencyLevel, write: ConsistencyLevel, peers: Vec<Peer>) -> Arc<ServerPool> {
1987        let pool = ServerPool::new(cfg(read, write), peers);
1988        pool.preselect_remote_racks();
1989        Arc::new(pool)
1990    }
1991
1992    #[test]
1993    fn local_node_only_short_circuits() {
1994        let p = pool(
1995            ConsistencyLevel::DcOne,
1996            ConsistencyLevel::DcOne,
1997            vec![peer(0, "dc1", "rA", 10, true, true)],
1998        );
1999        let mut req = Msg::new(1, MsgType::ReqRedisGet, true);
2000        req.set_routing(MsgRouting::LocalNodeOnly);
2001        assert_eq!(
2002            ClusterDispatcher::new(p).plan(&req, b"k"),
2003            DispatchPlan::LocalDatastore,
2004        );
2005    }
2006
2007    #[test]
2008    fn dc_one_read_targets_local_rack_when_present() {
2009        let p = pool(
2010            ConsistencyLevel::DcOne,
2011            ConsistencyLevel::DcOne,
2012            vec![
2013                peer(0, "dc1", "rA", 10, true, true),
2014                peer(1, "dc1", "rB", 20, false, true),
2015                peer(2, "dc2", "rA", 30, false, false),
2016            ],
2017        );
2018        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2019        // Any key resolves to peer 0 in rack rA (single-token continuum).
2020        let plan = ClusterDispatcher::new(p).plan(&req, b"hello");
2021        assert!(matches!(plan, DispatchPlan::LocalDatastore));
2022    }
2023
2024    /// Regression: a DC_ONE WRITE must fan out to every local-DC
2025    /// replica (each rack is a replica), not just the rack-local node.
2026    /// C Dynomite replicates a DC_ONE write within the local DC; the
2027    /// Rust port previously picked one target for both reads and
2028    /// writes, so a write landed on a single backend and never
2029    /// replicated -- reads from other replicas then missed it (found
2030    /// by the EC2 differential vs C).
2031    #[test]
2032    fn dc_one_write_fans_out_to_all_local_replicas() {
2033        // Three local-DC replicas, one per rack (full-replica layout,
2034        // n_val = 3). Every node owns the whole ring (token 0), so a
2035        // write must reach all three.
2036        let p = pool(
2037            ConsistencyLevel::DcOne,
2038            ConsistencyLevel::DcOne,
2039            vec![
2040                peer(0, "dc1", "rA", 0, true, true),
2041                peer(1, "dc1", "rB", 0, false, true),
2042                peer(2, "dc1", "rC", 0, false, true),
2043            ],
2044        );
2045        // A write (is_read = false): ReqRedisSet.
2046        let mut req = Msg::new(1, MsgType::ReqRedisSet, false);
2047        req.flags_mut().is_read = false;
2048        req.push_key(crate::msg::keypos::KeyPos::without_tag(b"k".to_vec()));
2049        match ClusterDispatcher::new(p).plan(&req, b"k") {
2050            DispatchPlan::Replicas { targets, .. } => {
2051                assert_eq!(
2052                    targets.len(),
2053                    3,
2054                    "DC_ONE write fans out to all 3 local-DC replicas"
2055                );
2056                for t in &targets {
2057                    assert_eq!(t.dc, "dc1");
2058                }
2059            }
2060            other => panic!("expected a 3-replica fan-out for a DC_ONE write, got {other:?}"),
2061        }
2062    }
2063
2064    /// Regression: a DC_ONE write on a single-node pool (the only
2065    /// local-DC replica is this node) must resolve to `LocalDatastore`,
2066    /// not a one-element `Replicas` fan-out. The fan-out form requires a
2067    /// wired backend channel; the embedded wire server serves local
2068    /// traffic through the local-datastore hook via `LocalDatastore`, so
2069    /// a lone-replica write that planned as `Replicas` returned a
2070    /// no-quorum error over the wire. The multi-replica fan-out above is
2071    /// unaffected. See commit 45ba694 (the DC_ONE write-fan-out change
2072    /// that introduced the regression) and the embed_api wire test
2073    /// `embedded_listen_socket_serves_a_wire_client`.
2074    #[test]
2075    fn dc_one_write_on_lone_replica_serves_local() {
2076        let p = pool(
2077            ConsistencyLevel::DcOne,
2078            ConsistencyLevel::DcOne,
2079            vec![peer(0, "dc1", "rA", 0, true, true)],
2080        );
2081        let mut req = Msg::new(1, MsgType::ReqRedisSet, false);
2082        req.flags_mut().is_read = false;
2083        req.push_key(crate::msg::keypos::KeyPos::without_tag(b"k".to_vec()));
2084        assert!(
2085            matches!(
2086                ClusterDispatcher::new(p).plan(&req, b"k"),
2087                DispatchPlan::LocalDatastore
2088            ),
2089            "a DC_ONE write whose only local replica is this node must \
2090             serve LocalDatastore, not fan out as a 1-element Replicas plan"
2091        );
2092    }
2093
2094    /// A DC_ONE READ still picks the single rack-local replica
2095    /// (LocalDatastore here since the local node is a replica).
2096    #[test]
2097    fn dc_one_read_picks_one_replica() {
2098        let p = pool(
2099            ConsistencyLevel::DcOne,
2100            ConsistencyLevel::DcOne,
2101            vec![
2102                peer(0, "dc1", "rA", 0, true, true),
2103                peer(1, "dc1", "rB", 0, false, true),
2104                peer(2, "dc1", "rC", 0, false, true),
2105            ],
2106        );
2107        let mut req = Msg::new(1, MsgType::ReqRedisGet, true);
2108        req.push_key(crate::msg::keypos::KeyPos::without_tag(b"k".to_vec()));
2109        // Local node is a replica -> served locally, not fanned out.
2110        assert!(matches!(
2111            ClusterDispatcher::new(p).plan(&req, b"k"),
2112            DispatchPlan::LocalDatastore
2113        ));
2114    }
2115
2116    #[test]
2117    fn dc_one_partitions_the_ring_within_a_shared_rack() {
2118        // The deployment topology that makes DC_ONE route by key: all
2119        // same-DC nodes share ONE rack (rA) and their tokens partition
2120        // the ring. A key whose token falls in a REMOTE node's range
2121        // must route to that node, not collapse to the local node.
2122        // This is the multi-region EC2 model (per-DC token partition
2123        // in a shared rack); one-node-per-rack instead makes every
2124        // rack own the whole ring and every DC_ONE read stay local.
2125        let p = pool(
2126            ConsistencyLevel::DcOne,
2127            ConsistencyLevel::DcOne,
2128            vec![
2129                peer(0, "dc1", "rA", 0, true, true),
2130                peer(1, "dc1", "rA", 1_431_655_765, false, true),
2131                peer(2, "dc1", "rA", 2_863_311_530, false, true),
2132            ],
2133        );
2134        let disp = ClusterDispatcher::new(p);
2135        // Probe many keys; at least one must route to a REMOTE peer
2136        // (not LocalDatastore), proving the ring is partitioned across
2137        // the shared-rack nodes rather than always-local.
2138        let mut saw_remote = false;
2139        let mut saw_local = false;
2140        for i in 0..200u32 {
2141            let key = format!("key-{i}");
2142            let req = Msg::new(1, MsgType::ReqRedisGet, true);
2143            match disp.plan(&req, key.as_bytes()) {
2144                DispatchPlan::LocalDatastore => saw_local = true,
2145                DispatchPlan::Replicas { targets, .. } => {
2146                    assert_eq!(targets.len(), 1, "DC_ONE picks one owner");
2147                    assert_eq!(targets[0].dc, "dc1");
2148                    saw_remote = true;
2149                }
2150                other => panic!("unexpected plan {other:?}"),
2151            }
2152        }
2153        assert!(
2154            saw_remote,
2155            "keys owned by a remote same-rack node must route to it, not stay local"
2156        );
2157        assert!(saw_local, "keys owned by the local node stay local");
2158    }
2159
2160    /// Regression: a request forwarded to a REMOTE peer must be
2161    /// tagged `DmsgType::ReqForward`, not `Req`. The receiver only
2162    /// hands a `ReqForward` straight to its local datastore; a plain
2163    /// `Req` makes it re-route (re-hash / re-plan), so a write owned
2164    /// by a remote node never lands in that node's backend and is
2165    /// lost. Found on the EC2 cluster: the token owner received zero
2166    /// forwarded writes.
2167    #[tokio::test]
2168    async fn remote_forward_is_tagged_req_forward() {
2169        // Two same-rack peers with u32-spread tokens; peer 0 is
2170        // local. Find a key that routes to the remote peer 1 and
2171        // assert the captured OutboundRequest carries ReqForward.
2172        let p = pool(
2173            ConsistencyLevel::DcOne,
2174            ConsistencyLevel::DcOne,
2175            vec![
2176                peer(0, "dc1", "rA", 0, true, true),
2177                peer(1, "dc1", "rA", 2_147_483_648, false, true),
2178            ],
2179        );
2180        let (tx, mut rx) = mpsc::channel::<crate::net::server::OutboundRequest>(64);
2181        let disp = ClusterDispatcher::new(p).with_peer_backend(1, tx);
2182        // Drive keys until one routes to the remote peer.
2183        let pool_buf = crate::io::mbuf::MbufPool::default();
2184        let mut forwarded_ty = None;
2185        for i in 0..500u32 {
2186            let key = format!("k{i}");
2187            let mut req = Msg::new(1, MsgType::ReqRedisSet, false);
2188            req.push_key(crate::msg::keypos::KeyPos::without_tag(key.into_bytes()));
2189            let mut buf = pool_buf.get();
2190            buf.copy_from_slice(b"*3\r\n$3\r\nSET\r\n$1\r\nx\r\n$1\r\ny\r\n");
2191            req.mbufs_mut().push_back(buf);
2192            let (resp_tx, _resp_rx) = mpsc::channel(1);
2193            let _ = disp.dispatch(req, resp_tx);
2194            if let Ok(env) = rx.try_recv() {
2195                forwarded_ty = Some(env.ty);
2196                break;
2197            }
2198        }
2199        assert_eq!(
2200            forwarded_ty,
2201            Some(crate::proto::dnode::DmsgType::ReqForward),
2202            "a forward to a remote peer must be ReqForward so the target serves it locally"
2203        );
2204    }
2205
2206    #[test]
2207    fn dc_quorum_fans_out_local_dc() {
2208        let p = pool(
2209            ConsistencyLevel::DcQuorum,
2210            ConsistencyLevel::DcQuorum,
2211            vec![
2212                peer(0, "dc1", "rA", 10, true, true),
2213                peer(1, "dc1", "rB", 20, false, true),
2214                peer(2, "dc2", "rA", 30, false, false),
2215            ],
2216        );
2217        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2218        let plan = ClusterDispatcher::new(p).plan(&req, b"k");
2219        match plan {
2220            DispatchPlan::Replicas { targets: rs, .. } => {
2221                assert_eq!(rs.len(), 2);
2222                for r in rs {
2223                    assert_eq!(r.dc, "dc1");
2224                }
2225            }
2226            _ => panic!("expected replicas"),
2227        }
2228    }
2229
2230    #[test]
2231    fn dc_each_safe_quorum_fans_out_per_dc() {
2232        let p = pool(
2233            ConsistencyLevel::DcEachSafeQuorum,
2234            ConsistencyLevel::DcEachSafeQuorum,
2235            vec![
2236                peer(0, "dc1", "rA", 10, true, true),
2237                peer(1, "dc2", "rA", 20, false, false),
2238            ],
2239        );
2240        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2241        let plan = ClusterDispatcher::new(p).plan(&req, b"k");
2242        match plan {
2243            DispatchPlan::Replicas { targets: rs, .. } => {
2244                assert_eq!(rs.len(), 2);
2245                let dcs: Vec<&str> = rs.iter().map(|r| r.dc.as_str()).collect();
2246                assert!(dcs.contains(&"dc1"));
2247                assert!(dcs.contains(&"dc2"));
2248            }
2249            _ => panic!("expected replicas"),
2250        }
2251    }
2252
2253    #[test]
2254    fn no_routable_peers_returns_no_targets() {
2255        let mut p0 = peer(0, "dc1", "rA", 10, true, true);
2256        p0.set_state(PeerState::Down, 0);
2257        let p = pool(
2258            ConsistencyLevel::DcQuorum,
2259            ConsistencyLevel::DcQuorum,
2260            vec![p0],
2261        );
2262        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2263        let plan = ClusterDispatcher::new(p).plan(&req, b"k");
2264        assert_eq!(plan, DispatchPlan::NoTargets);
2265    }
2266
2267    /// Regression: any code path that returns
2268    /// `DispatchOutcome::Error` used to send 0 wire bytes to the
2269    /// client because [`crate::msg::response::make_error`] never
2270    /// attached the wire-format error string. The client then
2271    /// hung until its read timeout. After the fix, the error
2272    /// response carries a parseable `-Dynomite: ...` reply that
2273    /// the client can render as an error.
2274    #[test]
2275    fn no_targets_error_response_carries_dynomite_wire_bytes() {
2276        let mut p0 = peer(0, "dc1", "rA", 10, true, true);
2277        p0.set_state(PeerState::Down, 0);
2278        let p = pool(
2279            ConsistencyLevel::DcQuorum,
2280            ConsistencyLevel::DcQuorum,
2281            vec![p0],
2282        );
2283        let disp = ClusterDispatcher::new(p);
2284        let mut req = Msg::new(1, MsgType::ReqRedisGet, true);
2285        req.push_key(crate::msg::keypos::KeyPos::without_tag(b"k".to_vec()));
2286        let (tx, _rx) = mpsc::channel(1);
2287        let outcome = disp.dispatch(req, tx);
2288        match outcome {
2289            DispatchOutcome::Error(rsp) => {
2290                assert_eq!(rsp.ty(), MsgType::RspRedisError);
2291                assert!(rsp.flags().is_error);
2292                let bytes: Vec<u8> = rsp
2293                    .mbufs()
2294                    .iter()
2295                    .flat_map(|b| b.readable().to_vec())
2296                    .collect();
2297                assert!(
2298                    !bytes.is_empty(),
2299                    "NoTargets must produce on-wire bytes, not a 0-byte hang"
2300                );
2301                assert!(bytes.starts_with(b"-Dynomite: "));
2302                assert!(bytes.ends_with(b"\r\n"));
2303                assert_eq!(rsp.mlen() as usize, bytes.len());
2304            }
2305            other => panic!("expected DispatchOutcome::Error, got {other:?}"),
2306        }
2307    }
2308
2309    /// Memcache traffic with no quorum-eligible target must
2310    /// surface a `SERVER_ERROR ...\r\n` reply rather than
2311    /// hanging the client.
2312    #[test]
2313    fn no_targets_error_response_memcache_wire_bytes() {
2314        // Build a memcache pool so the dispatcher's err_type
2315        // selection lands on `RspMcServerError` (the dispatcher
2316        // currently keys off the request `MsgType`, so a
2317        // memcache request flows through the memcache wire
2318        // shape).
2319        let mut cfg = cfg(ConsistencyLevel::DcQuorum, ConsistencyLevel::DcQuorum);
2320        cfg.data_store = DataStore::Memcache;
2321        let mut p0 = peer(0, "dc1", "rA", 10, true, true);
2322        p0.set_state(PeerState::Down, 0);
2323        let pool_arc = ServerPool::new(cfg, vec![p0]);
2324        pool_arc.preselect_remote_racks();
2325        let disp = ClusterDispatcher::new(Arc::new(pool_arc));
2326        let mut req = Msg::new(1, MsgType::ReqMcGet, true);
2327        req.push_key(crate::msg::keypos::KeyPos::without_tag(b"k".to_vec()));
2328        let (tx, _rx) = mpsc::channel(1);
2329        let outcome = disp.dispatch(req, tx);
2330        match outcome {
2331            DispatchOutcome::Error(rsp) => {
2332                assert_eq!(rsp.ty(), MsgType::RspMcServerError);
2333                let bytes: Vec<u8> = rsp
2334                    .mbufs()
2335                    .iter()
2336                    .flat_map(|b| b.readable().to_vec())
2337                    .collect();
2338                assert!(
2339                    !bytes.is_empty(),
2340                    "NoTargets must produce on-wire bytes, not a 0-byte hang"
2341                );
2342                assert!(bytes.starts_with(b"SERVER_ERROR "));
2343                assert!(bytes.ends_with(b"\r\n"));
2344            }
2345            other => panic!("expected DispatchOutcome::Error, got {other:?}"),
2346        }
2347    }
2348
2349    use crate::cluster::pool::{BucketType, PoolConfig};
2350
2351    fn pool_with_bucket_types(
2352        pool_read: ConsistencyLevel,
2353        pool_write: ConsistencyLevel,
2354        bucket_types: Vec<BucketType>,
2355        default_bucket_type: Option<&str>,
2356        peers: Vec<Peer>,
2357    ) -> Arc<ServerPool> {
2358        let cfg = PoolConfig {
2359            read_consistency: pool_read,
2360            write_consistency: pool_write,
2361            dc: "dc1".into(),
2362            rack: "rA".into(),
2363            bucket_types,
2364            default_bucket_type: default_bucket_type.map(str::to_string),
2365            ..PoolConfig::default()
2366        };
2367        let pool = ServerPool::new(cfg, peers);
2368        pool.preselect_remote_racks();
2369        Arc::new(pool)
2370    }
2371
2372    fn three_local_peers() -> Vec<Peer> {
2373        vec![
2374            peer(0, "dc1", "rA", 10, true, true),
2375            peer(1, "dc1", "rB", 20, false, true),
2376            peer(2, "dc1", "rC", 30, false, true),
2377        ]
2378    }
2379
2380    #[test]
2381    fn bucket_type_overrides_pool_consistency() {
2382        // Pool default is DC_ONE, the bucket forces DC_QUORUM.
2383        let bts = vec![BucketType {
2384            name: "hot".into(),
2385            read_consistency: ConsistencyLevel::DcQuorum,
2386            write_consistency: ConsistencyLevel::DcQuorum,
2387            n_val: 0,
2388        }];
2389        let p = pool_with_bucket_types(
2390            ConsistencyLevel::DcOne,
2391            ConsistencyLevel::DcOne,
2392            bts,
2393            None,
2394            three_local_peers(),
2395        );
2396        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2397        let plan = ClusterDispatcher::new(p).plan(&req, b"hot/key1");
2398        match plan {
2399            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 3),
2400            other => panic!("expected DC_QUORUM fan-out, got {other:?}"),
2401        }
2402    }
2403
2404    #[test]
2405    fn slashless_key_falls_back_to_pool_default() {
2406        let bts = vec![BucketType {
2407            name: "hot".into(),
2408            read_consistency: ConsistencyLevel::DcQuorum,
2409            write_consistency: ConsistencyLevel::DcQuorum,
2410            n_val: 0,
2411        }];
2412        let p = pool_with_bucket_types(
2413            ConsistencyLevel::DcOne,
2414            ConsistencyLevel::DcOne,
2415            bts,
2416            None,
2417            three_local_peers(),
2418        );
2419        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2420        let plan = ClusterDispatcher::new(p).plan(&req, b"plain-key");
2421        // No slash and no default_bucket_type -> pool DC_ONE.
2422        // The local rack hosts peer 0 so the plan short-circuits.
2423        assert!(matches!(plan, DispatchPlan::LocalDatastore));
2424    }
2425
2426    #[test]
2427    fn unknown_bucket_uses_default_bucket_type_when_set() {
2428        let bts = vec![BucketType {
2429            name: "safe".into(),
2430            read_consistency: ConsistencyLevel::DcQuorum,
2431            write_consistency: ConsistencyLevel::DcQuorum,
2432            n_val: 0,
2433        }];
2434        let p = pool_with_bucket_types(
2435            ConsistencyLevel::DcOne,
2436            ConsistencyLevel::DcOne,
2437            bts,
2438            Some("safe"),
2439            three_local_peers(),
2440        );
2441        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2442        // Slashless key: bucket is None, default_bucket_type=safe applies
2443        // so we get the bucket-type's DC_QUORUM fan-out.
2444        let plan = ClusterDispatcher::new(p.clone()).plan(&req, b"plain-key");
2445        match plan {
2446            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 3),
2447            other => panic!("expected DC_QUORUM via default bucket, got {other:?}"),
2448        }
2449        // Slashed key with an unknown bucket prefix also falls
2450        // through to the default bucket type.
2451        let plan = ClusterDispatcher::new(p).plan(&req, b"unknown-bucket/key");
2452        match plan {
2453            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 3),
2454            other => panic!("expected DC_QUORUM via default bucket, got {other:?}"),
2455        }
2456    }
2457
2458    #[test]
2459    fn unknown_bucket_with_no_default_uses_pool_default() {
2460        let bts = vec![BucketType {
2461            name: "safe".into(),
2462            read_consistency: ConsistencyLevel::DcQuorum,
2463            write_consistency: ConsistencyLevel::DcQuorum,
2464            n_val: 0,
2465        }];
2466        let p = pool_with_bucket_types(
2467            ConsistencyLevel::DcOne,
2468            ConsistencyLevel::DcOne,
2469            bts,
2470            None,
2471            three_local_peers(),
2472        );
2473        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2474        let plan = ClusterDispatcher::new(p).plan(&req, b"unknown-bucket/key");
2475        assert!(matches!(plan, DispatchPlan::LocalDatastore));
2476    }
2477
2478    #[test]
2479    fn n_val_one_caps_replicas_to_first_target() {
2480        let bts = vec![BucketType {
2481            name: "thin".into(),
2482            read_consistency: ConsistencyLevel::DcQuorum,
2483            write_consistency: ConsistencyLevel::DcQuorum,
2484            n_val: 1,
2485        }];
2486        let p = pool_with_bucket_types(
2487            ConsistencyLevel::DcOne,
2488            ConsistencyLevel::DcOne,
2489            bts,
2490            None,
2491            three_local_peers(),
2492        );
2493        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2494        let plan = ClusterDispatcher::new(p).plan(&req, b"thin/key");
2495        match plan {
2496            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 1),
2497            other => panic!("expected single-target plan, got {other:?}"),
2498        }
2499    }
2500
2501    #[test]
2502    fn n_val_two_caps_replicas_to_first_two_targets() {
2503        let bts = vec![BucketType {
2504            name: "medium".into(),
2505            read_consistency: ConsistencyLevel::DcQuorum,
2506            write_consistency: ConsistencyLevel::DcQuorum,
2507            n_val: 2,
2508        }];
2509        let p = pool_with_bucket_types(
2510            ConsistencyLevel::DcOne,
2511            ConsistencyLevel::DcOne,
2512            bts,
2513            None,
2514            three_local_peers(),
2515        );
2516        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2517        let plan = ClusterDispatcher::new(p).plan(&req, b"medium/key");
2518        match plan {
2519            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 2),
2520            other => panic!("expected two-target plan, got {other:?}"),
2521        }
2522    }
2523
2524    #[test]
2525    fn n_val_zero_does_not_cap() {
2526        let bts = vec![BucketType {
2527            name: "any".into(),
2528            read_consistency: ConsistencyLevel::DcQuorum,
2529            write_consistency: ConsistencyLevel::DcQuorum,
2530            n_val: 0,
2531        }];
2532        let p = pool_with_bucket_types(
2533            ConsistencyLevel::DcOne,
2534            ConsistencyLevel::DcOne,
2535            bts,
2536            None,
2537            three_local_peers(),
2538        );
2539        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2540        let plan = ClusterDispatcher::new(p).plan(&req, b"any/key");
2541        match plan {
2542            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 3),
2543            other => panic!("expected uncapped plan, got {other:?}"),
2544        }
2545    }
2546
2547    #[test]
2548    fn n_val_larger_than_replicas_is_a_no_op() {
2549        let bts = vec![BucketType {
2550            name: "big".into(),
2551            read_consistency: ConsistencyLevel::DcQuorum,
2552            write_consistency: ConsistencyLevel::DcQuorum,
2553            n_val: 7,
2554        }];
2555        let p = pool_with_bucket_types(
2556            ConsistencyLevel::DcOne,
2557            ConsistencyLevel::DcOne,
2558            bts,
2559            None,
2560            three_local_peers(),
2561        );
2562        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2563        let plan = ClusterDispatcher::new(p).plan(&req, b"big/key");
2564        match plan {
2565            DispatchPlan::Replicas { targets: rs, .. } => assert_eq!(rs.len(), 3),
2566            other => panic!("expected uncapped plan, got {other:?}"),
2567        }
2568    }
2569
2570    /// Smoke test for the failure-cause counter wiring: a
2571    /// `NoTargets` plan increments the labelled counter
2572    /// exactly once.
2573    #[test]
2574    fn no_targets_records_failure_metric() {
2575        let mut p0 = peer(0, "dc1", "rA", 10, true, true);
2576        p0.set_state(PeerState::Down, 0);
2577        let p = pool(
2578            ConsistencyLevel::DcQuorum,
2579            ConsistencyLevel::DcQuorum,
2580            vec![p0],
2581        );
2582        let metrics = Arc::new(crate::stats::FailureMetrics::new());
2583        let disp = ClusterDispatcher::new(p).with_failure_metrics(metrics.clone());
2584        let req = Msg::new(1, MsgType::ReqRedisGet, true);
2585        assert_eq!(disp.plan(&req, b"k"), DispatchPlan::NoTargets);
2586        let snap = metrics.snapshot();
2587        assert_eq!(snap.no_targets.len(), 1);
2588        let entry = &snap.no_targets[0];
2589        assert_eq!(entry.dc, "dc1");
2590        assert_eq!(entry.rack, "rA");
2591        assert_eq!(entry.consistency, ConsistencyLevel::DcQuorum);
2592        assert_eq!(entry.count, 1);
2593    }
2594
2595    /// A closed mpsc channel returns `Closed` from
2596    /// `try_send`. Wire the dispatcher's local-datastore path
2597    /// to such a channel, fire one request, and assert the
2598    /// `dispatch_backend_send_closed_total` counter ticks by
2599    /// exactly one.
2600    #[tokio::test]
2601    async fn closed_backend_channel_records_closed_metric() {
2602        let p = pool(
2603            ConsistencyLevel::DcOne,
2604            ConsistencyLevel::DcOne,
2605            vec![peer(0, "dc1", "rA", 10, true, true)],
2606        );
2607        let (tx, rx) = mpsc::channel::<crate::net::server::OutboundRequest>(4);
2608        drop(rx);
2609        let metrics = Arc::new(crate::stats::FailureMetrics::new());
2610        let disp = ClusterDispatcher::new(p)
2611            .with_backend(tx)
2612            .with_failure_metrics(metrics.clone());
2613        let mut req = Msg::new(1, MsgType::ReqRedisGet, true);
2614        // Attach a non-empty mbuf so the dispatcher actually
2615        // attempts the try_send (empty bytes short-circuit
2616        // to Drop before the channel is touched).
2617        let pool_buf = crate::io::mbuf::MbufPool::default();
2618        let mut buf = pool_buf.get();
2619        buf.copy_from_slice(b"PING\r\n");
2620        req.mbufs_mut().push_back(buf);
2621        let (resp_tx, _resp_rx) = mpsc::channel(1);
2622        let outcome = disp.dispatch(req, resp_tx);
2623        assert!(matches!(outcome, DispatchOutcome::Error(_)));
2624        let snap = metrics.snapshot();
2625        assert_eq!(snap.backend_send_closed, 1);
2626        assert_eq!(snap.backend_send_full, 0);
2627    }
2628
2629    /// Integration-style unit test: build a two-peer pool,
2630    /// mark one peer Down, drive 100 dispatches across the
2631    /// ring, and assert the `dispatch_no_targets_total`
2632    /// counter reflects every observed `NoTargets` plan.
2633    #[test]
2634    fn two_peer_pool_with_one_down_records_per_key_no_targets() {
2635        let cfg = crate::cluster::PoolConfig {
2636            dc: "dc1".into(),
2637            rack: "rA".into(),
2638            read_consistency: ConsistencyLevel::DcQuorum,
2639            write_consistency: ConsistencyLevel::DcQuorum,
2640            ..crate::cluster::PoolConfig::default()
2641        };
2642        // Single-rack two-peer ring: peer 0 owns the upper
2643        // half via wrap-around (token 2_147_483_648), peer 1
2644        // owns the lower half (token 0 plus the boundary up to
2645        // 2_147_483_648). With both peers in rack `rA` the
2646        // continuum has two entries, so each key maps to
2647        // exactly one peer; marking peer 1 Down causes its
2648        // arc to produce `NoTargets`.
2649        let p0 = peer(0, "dc1", "rA", 2_147_483_648, true, true);
2650        let mut p1 = peer(1, "dc1", "rA", 0, false, true);
2651        p1.set_state(PeerState::Down, 0);
2652        let pool_arc = ServerPool::new(cfg, vec![p0, p1]);
2653        pool_arc.preselect_remote_racks();
2654        let metrics = Arc::new(crate::stats::FailureMetrics::new());
2655        let disp = ClusterDispatcher::new(Arc::new(pool_arc)).with_failure_metrics(metrics.clone());
2656        let mut planned_no_targets = 0u64;
2657        let mut planned_routable = 0u64;
2658        for i in 0..100u32 {
2659            let key = format!("k{i:03}");
2660            let req = Msg::new(u64::from(i), MsgType::ReqRedisGet, true);
2661            match disp.plan(&req, key.as_bytes()) {
2662                DispatchPlan::NoTargets => planned_no_targets += 1,
2663                DispatchPlan::Replicas { .. } | DispatchPlan::LocalDatastore => {
2664                    planned_routable += 1;
2665                }
2666                DispatchPlan::Drop => panic!("unexpected Drop in plan"),
2667            }
2668        }
2669        assert!(planned_no_targets > 0, "expected some NoTargets dispatches");
2670        assert!(planned_routable > 0, "expected some routable dispatches");
2671        let snap = metrics.snapshot();
2672        let counter_total: u64 = snap.no_targets.iter().map(|e| e.count).sum();
2673        assert_eq!(
2674            counter_total, planned_no_targets,
2675            "dispatch_no_targets_total must match observed NoTargets count",
2676        );
2677        // No `Closed`/`Full` channel errors expected: we did
2678        // not wire any backends.
2679        assert_eq!(snap.backend_send_full, 0);
2680        assert_eq!(snap.backend_send_closed, 0);
2681        assert!(snap.peer_send_full.is_empty());
2682        assert!(snap.peer_send_closed.is_empty());
2683    }
2684}