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