Skip to main content

dynomite/embed/
server.rs

1//! Embedding-API server runtime.
2//!
3//! [`Server`] owns the runtime data structures (cluster pool,
4//! dispatcher, stats, event bus, hooks). [`Server::start`] spawns
5//! the background tasks (gossip, stats aggregator, optional TCP
6//! listeners) on the current tokio runtime and returns a
7//! [`ServerHandle`].
8//!
9//! The handle exposes the public control surface documented in
10//! `docs/book/src/embedding/server.md`. It is `Clone + Send +
11//! Sync`; multiple consumers may hold one. Dropping the last
12//! handle does not shut the server down - only
13//! [`ServerHandle::shutdown`] does.
14
15use std::collections::HashMap;
16use std::net::SocketAddr;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::{Arc, OnceLock, Weak};
19
20use parking_lot::Mutex;
21use tokio::net::TcpListener;
22use tokio::sync::mpsc;
23use tokio::task::JoinHandle;
24use tokio::time::{Duration, MissedTickBehavior};
25use tokio_util::sync::CancellationToken;
26
27use crate::cluster::dispatch::{ClusterDispatcher, DispatchPlan};
28use crate::cluster::peer::{Peer, PeerEndpoint, PeerState};
29use crate::cluster::pool::{PoolConfig, ServerPool};
30use crate::conf::{ConfDynSeed, ConfPool, Config, DataStore};
31use crate::embed::error::EmbedError;
32use crate::embed::events::{CloseReason, ConnRoleTag, EventBus, EventStream, ServerEvent};
33use crate::embed::hooks::{
34    CryptoProvider, Datastore, LoggingMetricsSink, MetricsSink, SeedsProvider,
35};
36use crate::embed::snapshots::{DatacenterSnapshot, PeerSnapshot, RingSnapshot};
37use crate::events::EventManager;
38use crate::hashkit::DynToken;
39use crate::io::reactor::{ConnRole, TcpTransport};
40use crate::msg::Msg;
41use crate::net::client::client_loop;
42use crate::net::{ClientHandler, Conn};
43use crate::stats::{
44    describe_stats, MetricSpec, PoolField, PoolStats, ServerField, ServerStats, ServiceInfo,
45    Snapshot, Stats,
46};
47
48/// Bag of optional hook overrides supplied to a builder.
49///
50/// All fields are private so adding new hooks (for example, a
51/// `Box<dyn TransportListener>` slot) is not a SemVer-breaking
52/// change. Construct a `ServerHooks` only through
53/// [`crate::embed::ServerBuilder`] setters; read it back through
54/// the accessor methods below.
55#[non_exhaustive]
56pub struct ServerHooks {
57    pub(crate) datastore: Option<Box<dyn Datastore>>,
58    pub(crate) seeds: Option<Box<dyn SeedsProvider>>,
59    pub(crate) crypto: Option<Box<dyn CryptoProvider>>,
60    pub(crate) metrics: Option<Box<dyn MetricsSink>>,
61}
62
63impl ServerHooks {
64    /// Borrow the configured [`Datastore`], if any.
65    #[must_use]
66    pub fn datastore(&self) -> Option<&dyn Datastore> {
67        self.datastore.as_deref()
68    }
69
70    /// Borrow the configured [`SeedsProvider`], if any.
71    #[must_use]
72    pub fn seeds(&self) -> Option<&dyn SeedsProvider> {
73        self.seeds.as_deref()
74    }
75
76    /// Borrow the configured [`CryptoProvider`], if any.
77    #[must_use]
78    pub fn crypto(&self) -> Option<&dyn CryptoProvider> {
79        self.crypto.as_deref()
80    }
81
82    /// Borrow the configured [`MetricsSink`], if any.
83    #[must_use]
84    pub fn metrics(&self) -> Option<&dyn MetricsSink> {
85        self.metrics.as_deref()
86    }
87}
88
89impl std::fmt::Debug for ServerHooks {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        f.debug_struct("ServerHooks")
92            .field("datastore", &self.datastore.is_some())
93            .field("seeds", &self.seeds.is_some())
94            .field("crypto", &self.crypto.is_some())
95            .field("metrics", &self.metrics.is_some())
96            .finish()
97    }
98}
99
100// -------- in-process registry -------------------------------------------
101
102/// Process-wide registry for in-process clusters.
103///
104/// When `inject_request` produces a `Replicas` plan that targets
105/// a peer co-located in the same process, the embed runtime
106/// looks up the target via this registry and forwards directly
107/// without the dnode wire path. This keeps the in-process tests
108/// from depending on real network I/O while preserving the
109/// production code path semantics (compute plan, deliver to
110/// target peers).
111fn registry() -> &'static Mutex<HashMap<SocketAddr, Weak<ServerInner>>> {
112    static R: OnceLock<Mutex<HashMap<SocketAddr, Weak<ServerInner>>>> = OnceLock::new();
113    R.get_or_init(|| Mutex::new(HashMap::new()))
114}
115
116fn registry_register(addr: SocketAddr, server: &Arc<ServerInner>) {
117    registry().lock().insert(addr, Arc::downgrade(server));
118}
119
120fn registry_remove(addr: SocketAddr) {
121    registry().lock().remove(&addr);
122}
123
124fn registry_lookup(addr: SocketAddr) -> Option<Arc<ServerInner>> {
125    registry().lock().get(&addr).and_then(Weak::upgrade)
126}
127
128// -------- inner state ----------------------------------------------------
129
130/// Generation counter for [`ServerEvent::ConfigReloaded`] /
131/// [`RingSnapshot`].
132fn next_generation() -> u64 {
133    static G: AtomicU64 = AtomicU64::new(0);
134    G.fetch_add(1, Ordering::Relaxed)
135}
136
137/// Connection-id allocator.
138fn next_conn_id() -> u64 {
139    static C: AtomicU64 = AtomicU64::new(0);
140    C.fetch_add(1, Ordering::Relaxed)
141}
142
143/// Inner state shared by every clone of a [`ServerHandle`].
144pub(crate) struct ServerInner {
145    pool: Arc<ServerPool>,
146    dispatcher: Arc<ClusterDispatcher>,
147    client_proto: DataStore,
148    stats: Arc<Stats>,
149    snapshot_cache: Arc<Mutex<Snapshot>>,
150    bus: EventBus,
151    events: Arc<EventManager>,
152    datastore: Arc<dyn Datastore>,
153    seeds: Box<dyn SeedsProvider>,
154    metrics: Box<dyn MetricsSink>,
155    crypto: Option<Box<dyn CryptoProvider>>,
156    listen_addr: Option<SocketAddr>,
157    dyn_listen_addr: Option<SocketAddr>,
158    cancel: CancellationToken,
159    pool_name: String,
160    config: Mutex<ConfPool>,
161    generation: AtomicU64,
162    command_extension: Option<Arc<dyn crate::embed::CommandExtension>>,
163}
164
165impl std::fmt::Debug for ServerInner {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        f.debug_struct("ServerInner")
168            .field("pool_name", &self.pool_name)
169            .field("listen", &self.listen_addr)
170            .field("dyn_listen", &self.dyn_listen_addr)
171            .finish_non_exhaustive()
172    }
173}
174
175impl ServerInner {
176    fn dispatch_local(&self, req: Msg) -> super::hooks::BoxFuture<'_, Result<Msg, EmbedError>> {
177        let fut = self.datastore.dispatch(req);
178        Box::pin(async move {
179            let rsp = fut.await.map_err(|e| EmbedError::Inject(e.to_string()))?;
180            Ok(rsp)
181        })
182    }
183}
184
185// -------- Server (configured-but-not-running) ----------------------------
186
187/// Configured-but-not-running server.
188///
189/// Build one with [`ServerBuilder`](crate::embed::ServerBuilder),
190/// then call [`Server::start`] to spawn the background tasks.
191#[derive(Debug)]
192pub struct Server {
193    pool_name: String,
194    pool: ConfPool,
195    cluster: Arc<ServerPool>,
196    hooks: ServerHooks,
197    stats: Arc<Stats>,
198    command_extension: Option<Arc<dyn crate::embed::CommandExtension>>,
199}
200
201impl Server {
202    /// Used internally by [`crate::embed::ServerBuilder::build`].
203    pub(crate) fn from_pool(
204        pool_name: String,
205        pool: ConfPool,
206        hooks: ServerHooks,
207        command_extension: Option<Arc<dyn crate::embed::CommandExtension>>,
208    ) -> Self {
209        let pool_cfg = PoolConfig::from_conf(&pool_name, &pool);
210        let local_peer = build_local_peer(&pool, &pool_cfg);
211        let mut peers = vec![local_peer];
212        if let Some(seeds) = pool.dyn_seeds.as_ref() {
213            let start = u32::try_from(peers.len()).unwrap_or(0);
214            peers.extend(peers_from_seeds(&pool_cfg, seeds, start));
215        }
216        let server_pool_arc = Arc::new(ServerPool::new(pool_cfg.clone(), peers));
217        server_pool_arc.preselect_remote_racks();
218
219        let stats = Arc::new(Stats::new(
220            ServiceInfo {
221                source: pool_cfg.name.clone(),
222                version: env!("CARGO_PKG_VERSION").to_string(),
223                rack: pool_cfg.rack.clone(),
224                dc: pool_cfg.dc.clone(),
225            },
226            PoolStats::new(&pool_cfg.name),
227            ServerStats::new("backend"),
228        ));
229
230        Self {
231            pool_name,
232            pool,
233            cluster: server_pool_arc,
234            hooks,
235            stats,
236            command_extension,
237        }
238    }
239
240    /// Borrow the [`crate::embed::CommandExtension`]
241    /// this server was built with, if any.
242    #[must_use]
243    pub fn command_extension(&self) -> Option<&Arc<dyn crate::embed::CommandExtension>> {
244        self.command_extension.as_ref()
245    }
246
247    /// The pool name configured in the YAML / builder.
248    #[must_use]
249    pub fn pool_name(&self) -> &str {
250        &self.pool_name
251    }
252
253    /// Spawn background tasks on the current tokio runtime and
254    /// return a [`ServerHandle`].
255    ///
256    /// `start` is non-blocking. The returned handle is `Clone +
257    /// Send + Sync`.
258    ///
259    /// # Client plane and in-process traffic
260    ///
261    /// A client connecting to the bound `listen:` socket is
262    /// served end to end: the request is parsed, routed through
263    /// the cluster dispatcher, and for a request that resolves to
264    /// the local node the configured [`Datastore`] hook produces
265    /// the reply -- the same hook
266    /// [`ServerHandle::inject_request`] uses for in-process
267    /// traffic. `inject_request` remains available for callers
268    /// that prefer to drive the engine without a socket hop.
269    ///
270    /// The `dyn_listen:` (peer-plane) socket binds for address
271    /// reporting; embedded multi-node setups forward between
272    /// in-process nodes through the in-process registry, so a raw
273    /// cross-process peer connection to `dyn_listen:` is accepted
274    /// and closed. Cross-process peer serving is provided by the
275    /// `dynomited` binary.
276    ///
277    /// # Examples
278    ///
279    /// ```no_run
280    /// use dynomite::embed::ServerBuilder;
281    /// use dynomite::conf::DataStore;
282    /// # tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap().block_on(async {
283    /// let server = ServerBuilder::new("dyn_o_mite")
284    ///     .listen("127.0.0.1:0".parse().unwrap())
285    ///     .dyn_listen("127.0.0.1:0".parse().unwrap())
286    ///     .data_store(DataStore::Valkey)
287    ///     .servers(vec![dynomite::conf::ConfServer::parse("127.0.0.1:6379:1").unwrap()])
288    ///     .tokens_str("0")
289    ///     .enable_gossip(false)
290    ///     .build()
291    ///     .unwrap();
292    /// let handle = server.start().await.unwrap();
293    /// handle.shutdown().await.unwrap();
294    /// # });
295    /// ```
296    pub async fn start(self) -> Result<ServerHandle, EmbedError> {
297        let Server {
298            pool_name,
299            pool,
300            cluster,
301            mut hooks,
302            stats,
303            command_extension,
304        } = self;
305
306        let datastore: Arc<dyn Datastore> = hooks
307            .datastore
308            .take()
309            .expect("invariant: builder always populates datastore")
310            .into();
311
312        let mut dispatcher = ClusterDispatcher::new(cluster.clone());
313        if let Some(ext) = command_extension.as_ref() {
314            dispatcher = dispatcher.with_command_extension(ext.clone());
315        }
316        // Serve client-plane requests that route to the local node
317        // through the same Datastore hook inject_request uses, so a
318        // connection accepted on the bound listen: socket is fully
319        // served (parse -> route -> reply) rather than closed.
320        dispatcher = dispatcher.with_local_datastore(Arc::clone(&datastore));
321        let dispatcher = Arc::new(dispatcher);
322        // Resolve the wire protocol the client proxy parses. The
323        // pool stores it as the numeric data_store selector;
324        // default to Valkey (RESP) when unset, matching the proxy
325        // default.
326        let client_proto = pool
327            .data_store
328            .and_then(|n| DataStore::from_int(n).ok())
329            .unwrap_or(DataStore::Valkey);
330        let bus = EventBus::new(64);
331        let events = Arc::new(EventManager::new(64));
332        let cancel = CancellationToken::new();
333        let snapshot_cache = Arc::new(Mutex::new(Snapshot::default()));
334
335        let seeds = hooks
336            .seeds
337            .take()
338            .expect("invariant: builder always populates seeds");
339        let metrics = hooks
340            .metrics
341            .take()
342            .unwrap_or_else(|| Box::new(LoggingMetricsSink::new(pool_name.clone())));
343
344        let transport = pool.transport.unwrap_or_default();
345        let (listen_listener, listen_addr) =
346            bind_listener(pool.listen.as_ref(), transport, &pool).await?;
347        let (dyn_listener, dyn_listen_addr) =
348            bind_listener(pool.dyn_listen.as_ref(), transport, &pool).await?;
349
350        let inner = Arc::new(ServerInner {
351            pool: cluster.clone(),
352            dispatcher,
353            client_proto,
354            stats: stats.clone(),
355            snapshot_cache: snapshot_cache.clone(),
356            bus: bus.clone(),
357            events: events.clone(),
358            datastore,
359            seeds,
360            metrics,
361            crypto: hooks.crypto,
362            listen_addr,
363            dyn_listen_addr,
364            cancel: cancel.clone(),
365            pool_name: pool_name.clone(),
366            config: Mutex::new(pool.clone()),
367            generation: AtomicU64::new(0),
368            command_extension,
369        });
370
371        // Register self in the in-process registry so peer
372        // forwarding can find this node.
373        if let Some(addr) = inner.dyn_listen_addr {
374            registry_register(addr, &inner);
375        }
376
377        // Spawn background tasks.
378        let mut tasks: Vec<JoinHandle<()>> = Vec::new();
379        tasks.push(tokio::spawn(stats_loop(inner.clone(), pool.stats_interval)));
380        tasks.push(tokio::spawn(metrics_loop(inner.clone())));
381        if pool.enable_gossip.unwrap_or(false) {
382            tasks.push(tokio::spawn(gossip_loop(
383                inner.clone(),
384                Duration::from_millis(
385                    u64::try_from(pool.gos_interval.unwrap_or(1_000)).unwrap_or(1_000),
386                ),
387            )));
388        }
389        if let (Some(listener), Some(addr)) = (listen_listener, listen_addr) {
390            tasks.push(tokio::spawn(accept_loop(
391                inner.clone(),
392                listener,
393                addr,
394                ConnRoleTag::Proxy,
395            )));
396        }
397        if let (Some(listener), Some(addr)) = (dyn_listener, dyn_listen_addr) {
398            tasks.push(tokio::spawn(accept_loop(
399                inner.clone(),
400                listener,
401                addr,
402                ConnRoleTag::DnodeProxy,
403            )));
404        }
405
406        Ok(ServerHandle {
407            inner,
408            tasks: Arc::new(Mutex::new(tasks)),
409        })
410    }
411}
412
413// -------- ServerHandle ---------------------------------------------------
414
415/// Cloneable handle to a running [`Server`].
416///
417/// The handle is the public control surface: shutdown, reload,
418/// stats, events, request injection, topology snapshots.
419#[derive(Clone)]
420pub struct ServerHandle {
421    inner: Arc<ServerInner>,
422    tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
423}
424
425impl std::fmt::Debug for ServerHandle {
426    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
427        f.debug_struct("ServerHandle")
428            .field("pool", &self.inner.pool_name)
429            .field("listen", &self.inner.listen_addr)
430            .field("dyn_listen", &self.inner.dyn_listen_addr)
431            .finish_non_exhaustive()
432    }
433}
434
435impl ServerHandle {
436    /// Borrow the configured [`CryptoProvider`], if one was
437    /// plugged. The default builder does not install one.
438    #[must_use]
439    pub fn crypto_provider(&self) -> Option<&dyn CryptoProvider> {
440        self.inner.crypto.as_deref()
441    }
442
443    /// Borrow the [`crate::embed::CommandExtension`] installed
444    /// at build time, if any.
445    #[must_use]
446    pub fn command_extension(&self) -> Option<Arc<dyn crate::embed::CommandExtension>> {
447        self.inner.command_extension.as_ref().map(Arc::clone)
448    }
449
450    /// Local listen address (post-bind), if a `listen:` was
451    /// configured.
452    #[must_use]
453    pub fn listen_addr(&self) -> Option<SocketAddr> {
454        self.inner.listen_addr
455    }
456
457    /// Local dnode listen address (post-bind), if a
458    /// `dyn_listen:` was configured.
459    #[must_use]
460    pub fn dyn_listen_addr(&self) -> Option<SocketAddr> {
461        self.inner.dyn_listen_addr
462    }
463
464    /// Subscribe to the [`ServerEvent`] broadcast.
465    ///
466    /// # Examples
467    ///
468    /// ```no_run
469    /// # use dynomite::embed::ServerBuilder;
470    /// # use dynomite::conf::DataStore;
471    /// # tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap().block_on(async {
472    /// let server = ServerBuilder::new("p")
473    ///     .listen("127.0.0.1:0".parse().unwrap())
474    ///     .dyn_listen("127.0.0.1:0".parse().unwrap())
475    ///     .data_store(DataStore::Valkey)
476    ///     .servers(vec![dynomite::conf::ConfServer::parse("127.0.0.1:6379:1").unwrap()])
477    ///     .tokens_str("0")
478    ///     .build().unwrap();
479    /// let handle = server.start().await.unwrap();
480    /// let _events = handle.subscribe_events();
481    /// handle.shutdown().await.unwrap();
482    /// # });
483    /// ```
484    #[must_use]
485    pub fn subscribe_events(&self) -> EventStream {
486        self.inner.bus.subscribe()
487    }
488
489    /// Borrow the [`EventManager`] handle that publishes
490    /// structured cluster-event payloads (`PeerUp`, `PeerDown`,
491    /// `GossipRoundComplete`, etc.).
492    ///
493    /// Returned as `Arc` so embedders can hold a long-lived
494    /// reference and pass it across crate boundaries (notably to
495    /// `dyniak`, which publishes AAE start/complete events).
496    #[must_use]
497    pub fn events(&self) -> Arc<EventManager> {
498        Arc::clone(&self.inner.events)
499    }
500
501    /// Latest stats snapshot.
502    #[must_use]
503    pub fn stats(&self) -> Snapshot {
504        self.inner.stats.snapshot()
505    }
506
507    /// Borrow the live [`Stats`] aggregator as a shared handle.
508    ///
509    /// Returns the same `Arc<Stats>` the runtime publishes to,
510    /// so an embedder can read counters as they update without
511    /// going through the periodic snapshot path. Useful for
512    /// in-process metrics integrations that prefer a pull model
513    /// (Prometheus scrape, OpenTelemetry pull readers) over the
514    /// push-on-flush path exposed by [`MetricsSink`].
515    ///
516    /// The returned handle is `Clone`; callers may share it
517    /// across threads. The aggregator outlives the handle for as
518    /// long as the [`ServerHandle`] tree is alive.
519    ///
520    /// [`MetricsSink`]: crate::embed::MetricsSink
521    #[must_use]
522    pub fn stats_handle(&self) -> Arc<Stats> {
523        Arc::clone(&self.inner.stats)
524    }
525
526    /// Manifest of every metric the engine emits.
527    #[must_use]
528    pub fn describe_stats(&self) -> Vec<MetricSpec> {
529        let _ = describe_stats(); // ensure the module is exercised
530        let mut out: Vec<MetricSpec> = Vec::new();
531        out.extend(crate::stats::POOL_CODEC.iter().copied());
532        out.extend(crate::stats::SERVER_CODEC.iter().copied());
533        out
534    }
535
536    /// Snapshot of every peer in the cluster.
537    #[must_use]
538    pub fn peers(&self) -> Vec<PeerSnapshot> {
539        self.inner
540            .pool
541            .peers()
542            .read()
543            .iter()
544            .map(PeerSnapshot::from)
545            .collect()
546    }
547
548    /// Snapshot of every datacenter and its racks.
549    #[must_use]
550    pub fn datacenters(&self) -> Vec<DatacenterSnapshot> {
551        self.inner
552            .pool
553            .datacenters()
554            .read()
555            .iter()
556            .map(DatacenterSnapshot::from)
557            .collect()
558    }
559
560    /// Snapshot of the token ring.
561    #[must_use]
562    pub fn ring(&self) -> RingSnapshot {
563        let dcs = self.inner.pool.datacenters().read();
564        let mut entries: Vec<(DynToken, u32)> = Vec::new();
565        for dc in dcs.iter() {
566            for rack in dc.racks() {
567                for c in rack.continuums() {
568                    entries.push((c.token.clone(), c.peer_idx));
569                }
570            }
571        }
572        RingSnapshot {
573            entries,
574            generation: self.inner.generation.load(Ordering::Relaxed),
575        }
576    }
577
578    /// Inject a parsed request as if it had arrived from a
579    /// client. Returns the response message.
580    ///
581    /// This is the test/embedding entry point that bypasses the
582    /// proxy listener. The dispatcher computes a routing plan; if
583    /// the plan resolves to the local datastore, the
584    /// [`Datastore::dispatch`]
585    /// hook is invoked. If the plan resolves to a remote peer
586    /// co-located in the same process (via the in-process
587    /// registry), the request is forwarded to that peer's
588    /// datastore and its response is returned.
589    pub async fn inject_request(&self, req: Msg) -> Result<Msg, EmbedError> {
590        // Use the parsed key (if any) for routing; if there are
591        // no keys, the dispatcher's `LocalDatastore` short-circuit
592        // applies.
593        let key: Vec<u8> = req
594            .keys()
595            .first()
596            .map(|kp| kp.tag_bytes().to_vec())
597            .unwrap_or_default();
598        let plan = self.inner.dispatcher.plan(&req, &key);
599        match plan {
600            DispatchPlan::Drop => {
601                // Drop: respond with an empty message of the same
602                // id so callers can correlate.
603                let mut rsp = Msg::new(req.id(), crate::msg::MsgType::Unknown, false);
604                rsp.set_parent_id(req.id());
605                Ok(rsp)
606            }
607            DispatchPlan::NoTargets => Err(EmbedError::Inject(
608                "cluster has no quorum-eligible targets".into(),
609            )),
610            DispatchPlan::LocalDatastore => self.inner.dispatch_local(req).await,
611            DispatchPlan::Replicas { targets, .. } => {
612                self.inner.stats.pool_incr_by(PoolField::ForwardError, 0); // touch counter table
613                self.inner.stats.server_incr(ServerField::ReadRequests);
614                // Snapshot the targets' addresses outside any
615                // lock so we don't hold the peer-list read guard
616                // across the awaits below.
617                let resolved: Vec<(bool, Option<SocketAddr>)> = {
618                    let peers = self.inner.pool.peers().read();
619                    targets
620                        .iter()
621                        .map(|t| {
622                            let peer = peers.get(t.peer_idx as usize);
623                            let is_local = peer.is_some_and(crate::cluster::peer::Peer::is_local);
624                            let addr = peer.and_then(|p| p.endpoint().pname().parse().ok());
625                            (is_local, addr)
626                        })
627                        .collect()
628                };
629                let mut last_err: Option<EmbedError> = None;
630                for (is_local, addr) in resolved {
631                    if is_local {
632                        return self.inner.dispatch_local(req).await;
633                    }
634                    if let Some(addr) = addr {
635                        if let Some(remote) = registry_lookup(addr) {
636                            let mut forwarded = Msg::new(req.id(), req.ty(), true);
637                            forwarded.set_parent_id(req.parent_id());
638                            match remote.datastore.dispatch(forwarded).await {
639                                Ok(rsp) => return Ok(rsp),
640                                Err(e) => last_err = Some(EmbedError::Inject(e.to_string())),
641                            }
642                        }
643                    }
644                }
645                if let Some(e) = last_err {
646                    return Err(e);
647                }
648                // No registered remote peer; fall back to the
649                // local datastore so the embed surface stays
650                // useful in mixed deployments.
651                self.inner.dispatch_local(req).await
652            }
653        }
654    }
655
656    /// Reload the configuration. Validates the new [`Config`]
657    /// before any state is touched; on failure the running config
658    /// is left untouched.
659    ///
660    /// This is the in-process equivalent of SIGHUP: instead of
661    /// re-reading the YAML from disk, the embedding API takes the
662    /// replacement `Config` directly.
663    pub async fn reload(&self, mut cfg: Config) -> Result<(), EmbedError> {
664        cfg.finalize();
665        cfg.validate()?;
666        let new_pool = cfg.pool().clone();
667        // Update the stored pool so subsequent .stats() / .peers()
668        // calls reflect the live config. The dispatcher's pool
669        // structure is left untouched here; a future stage will
670        // rebuild it under a write lock.
671        *self.inner.config.lock() = new_pool;
672        let gen_id = next_generation();
673        self.inner.generation.store(gen_id, Ordering::Relaxed);
674        self.inner
675            .bus
676            .send(ServerEvent::ConfigReloaded { generation: gen_id });
677        // Yield once so the broadcast send is observable on the
678        // calling task's scheduler before the function returns.
679        tokio::task::yield_now().await;
680        Ok(())
681    }
682
683    /// Graceful shutdown.
684    ///
685    /// Cancels every background task, deregisters from the
686    /// in-process peer registry, drains the join set, and
687    /// returns. Idempotent: calling `shutdown` after a previous
688    /// `shutdown` (or after [`ServerHandle::join`] returned)
689    /// completes successfully without doing any work.
690    pub async fn shutdown(&self) -> Result<(), EmbedError> {
691        self.inner.cancel.cancel();
692        if let Some(addr) = self.inner.dyn_listen_addr {
693            registry_remove(addr);
694        }
695        let drained: Vec<JoinHandle<()>> = std::mem::take(&mut *self.tasks.lock());
696        for t in drained {
697            // Tasks observe `cancel` and return promptly.
698            let _ = t.await;
699        }
700        Ok(())
701    }
702
703    /// Wait for every background task to complete.
704    ///
705    /// Unlike [`ServerHandle::shutdown`], `join` does not request
706    /// cancellation: it parks the caller until the runtime tasks
707    /// (gossip, stats aggregator, accept loops, metrics flusher)
708    /// finish on their own. The intended use is the
709    /// "start-then-park" pattern in long-running embedders that
710    /// drive shutdown from a separate signal handler:
711    ///
712    /// ```no_run
713    /// # use dynomite::embed::ServerBuilder;
714    /// # use dynomite::conf::DataStore;
715    /// # tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap().block_on(async {
716    /// let handle = ServerBuilder::new("p")
717    ///     .listen("127.0.0.1:0".parse().unwrap())
718    ///     .dyn_listen("127.0.0.1:0".parse().unwrap())
719    ///     .data_store(DataStore::Valkey)
720    ///     .servers(vec![dynomite::conf::ConfServer::parse("127.0.0.1:6379:1").unwrap()])
721    ///     .tokens_str("0")
722    ///     .build().unwrap()
723    ///     .start().await.unwrap();
724    /// let shutdown = handle.clone();
725    /// tokio::spawn(async move {
726    ///     // Replace with `tokio::signal::ctrl_c().await.unwrap();`
727    ///     // in a real embedder.
728    ///     shutdown.shutdown().await.unwrap();
729    /// });
730    /// handle.join().await;
731    /// # });
732    /// ```
733    ///
734    /// `join` resolves once every task spawned by
735    /// [`Server::start`] has returned. Calling `join` repeatedly
736    /// is safe; the second call returns immediately because the
737    /// task set is drained on the first.
738    pub async fn join(&self) {
739        let drained: Vec<JoinHandle<()>> = std::mem::take(&mut *self.tasks.lock());
740        for t in drained {
741            let _ = t.await;
742        }
743    }
744}
745
746// -------- helpers ---------------------------------------------------------
747
748fn build_local_peer(pool: &ConfPool, cfg: &PoolConfig) -> Peer {
749    let dyn_listen = pool.dyn_listen.as_ref().map_or_else(
750        || ("127.0.0.1".to_string(), 0u16),
751        |l| (l.name().to_string(), l.port()),
752    );
753    let tokens = pool
754        .tokens
755        .as_ref()
756        .map(|tl| {
757            tl.components()
758                .iter()
759                .map(|c| DynToken::from_u32(c.digits.parse::<u32>().unwrap_or(0)))
760                .collect::<Vec<_>>()
761        })
762        .unwrap_or_default();
763    let mut peer = Peer::new(
764        0,
765        PeerEndpoint::tcp(dyn_listen.0, dyn_listen.1),
766        cfg.rack.clone(),
767        cfg.dc.clone(),
768        if tokens.is_empty() {
769            vec![DynToken::from_u32(0)]
770        } else {
771            tokens
772        },
773        true,
774        true,
775        false,
776    );
777    let now_secs = std::time::SystemTime::now()
778        .duration_since(std::time::UNIX_EPOCH)
779        .map_or(0, |d| d.as_secs());
780    peer.set_state(PeerState::Normal, now_secs);
781    peer
782}
783
784fn peers_from_seeds(cfg: &PoolConfig, seeds: &[ConfDynSeed], start_idx: u32) -> Vec<Peer> {
785    let now_secs = std::time::SystemTime::now()
786        .duration_since(std::time::UNIX_EPOCH)
787        .map_or(0, |d| d.as_secs());
788    seeds
789        .iter()
790        .enumerate()
791        .map(|(i, s)| {
792            let tokens: Vec<DynToken> = s
793                .tokens()
794                .components()
795                .iter()
796                .map(|c| DynToken::from_u32(c.digits.parse::<u32>().unwrap_or(0)))
797                .collect();
798            let idx_off = u32::try_from(i).unwrap_or(0);
799            let mut p = Peer::new(
800                start_idx + idx_off,
801                PeerEndpoint::tcp(s.host().to_string(), s.port()),
802                s.rack().to_string(),
803                s.dc().to_string(),
804                if tokens.is_empty() {
805                    vec![DynToken::from_u32(0)]
806                } else {
807                    tokens
808                },
809                false,
810                s.dc() == cfg.dc,
811                false,
812            );
813            p.set_state(PeerState::Normal, now_secs);
814            p
815        })
816        .collect()
817}
818
819// Unused legacy helpers removed.
820
821async fn bind_listener(
822    listen: Option<&crate::conf::ConfListen>,
823    transport: crate::conf::Transport,
824    #[cfg_attr(not(feature = "quic"), allow(unused_variables))] pool: &crate::conf::ConfPool,
825) -> Result<(Option<TcpListener>, Option<SocketAddr>), EmbedError> {
826    let Some(l) = listen else {
827        return Ok((None, None));
828    };
829    let host = l.name();
830    let port = l.port();
831    if host.is_empty() {
832        return Ok((None, None));
833    }
834    let addr_str = format!("{host}:{port}");
835    let Ok(_addr) = addr_str.parse::<SocketAddr>() else {
836        return Ok((None, None));
837    };
838    match transport {
839        crate::conf::Transport::Tcp => {
840            let listener = TcpListener::bind(&addr_str).await?;
841            let local = listener.local_addr()?;
842            Ok((Some(listener), Some(local)))
843        }
844        #[cfg(feature = "quic")]
845        crate::conf::Transport::Quic => {
846            // The embedded server's accept loop closes every
847            // accepted client (see `accept_loop` below): it logs
848            // and closes. We still bind the QUIC
849            // listener up-front so the embedder sees a clean
850            // bind error if the cert / key paths are bad,
851            // matching the TCP path's eager-bind contract;
852            // dropping the listener here is fine because
853            // production QUIC traffic flows through the
854            // `dynomited` binary, not the embedded harness.
855            let cert = pool
856                .quic_cert_file
857                .as_ref()
858                .and_then(|p| p.to_str().map(str::to_owned))
859                .ok_or_else(|| {
860                    EmbedError::Build(
861                        "transport: quic requires quic_cert_file (UTF-8 path)".to_string(),
862                    )
863                })?;
864            let key = pool
865                .quic_key_file
866                .as_ref()
867                .and_then(|p| p.to_str().map(str::to_owned))
868                .ok_or_else(|| {
869                    EmbedError::Build(
870                        "transport: quic requires quic_key_file (UTF-8 path)".to_string(),
871                    )
872                })?;
873            let cfg = crate::net::QuicConfig::server_with_cert_paths(cert, key);
874            let parsed: SocketAddr = addr_str.parse().map_err(|e: std::net::AddrParseError| {
875                EmbedError::Build(format!("invalid quic listen address {addr_str}: {e}"))
876            })?;
877            let listener = crate::net::QuicListener::bind(parsed, cfg).await?;
878            let local = listener.local_addr();
879            // The QUIC listener is dropped on this scope exit;
880            // the in-process embedder does not drive a
881            // QUIC accept loop. Return `None` for the listener
882            // so the existing TCP-flavoured accept_loop is
883            // skipped, but report the local address for
884            // observability.
885            drop(listener);
886            Ok((None, Some(local)))
887        }
888        #[cfg(not(feature = "quic"))]
889        crate::conf::Transport::Quic => Err(EmbedError::Build(
890            "transport: quic requires the engine's `quic` Cargo feature".to_string(),
891        )),
892    }
893}
894
895// removed unused try_bind helper
896
897async fn accept_loop(
898    inner: Arc<ServerInner>,
899    listener: TcpListener,
900    addr: SocketAddr,
901    role: ConnRoleTag,
902) {
903    loop {
904        tokio::select! {
905            biased;
906            () = inner.cancel.cancelled() => return,
907            res = listener.accept() => {
908                let Ok((sock, peer)) = res else { return };
909                let conn_id = next_conn_id();
910                inner.bus.send(ServerEvent::ConnectionAccepted {
911                    conn_id,
912                    role,
913                    local_addr: Some(addr),
914                });
915                if matches!(role, ConnRoleTag::Proxy) {
916                    // Client plane: serve the connection through the
917                    // same per-client pipeline the standalone proxy
918                    // uses (parse -> route via the dispatcher ->
919                    // write the reply). This is the wire entry point
920                    // for an embedded node; `inject_request` remains
921                    // available for in-process traffic.
922                    //
923                    // Small Redis/memcache requests should not
924                    // wait on Nagle; a peer that closed before
925                    // the option applies is harmless.
926                    let _ = sock.set_nodelay(true);
927                    let dispatcher: Arc<dyn crate::net::dispatcher::Dispatcher> =
928                        inner.dispatcher.clone();
929                    let client_proto = inner.client_proto;
930                    let bus = inner.bus.clone();
931                    tokio::spawn(async move {
932                        let transport = Box::new(TcpTransport::new(sock, ConnRole::Client));
933                        let conn = Conn::new(transport, ConnRole::Client);
934                        let (tx, rx) = mpsc::channel(64);
935                        let handler = ClientHandler::new(dispatcher, tx, client_proto);
936                        let outcome = client_loop(conn, handler, rx).await;
937                        let close_reason = match outcome {
938                            Ok(()) => CloseReason::PeerEof,
939                            Err(_) => CloseReason::IoError,
940                        };
941                        bus.send(ServerEvent::ConnectionClosed {
942                            conn_id,
943                            reason: close_reason,
944                        });
945                    });
946                } else {
947                    // Peer plane: cross-process dnode serving is
948                    // driven by the `dynomited` binary. Embedded
949                    // multi-node setups forward between in-process
950                    // nodes through the in-process registry (see
951                    // `registry_register`), so the embedded
952                    // dyn_listen socket is bound for address
953                    // reporting and accepts then closes a raw
954                    // cross-process peer connection.
955                    tracing::debug!(
956                        listen = %addr,
957                        peer = %peer,
958                        role = ?role,
959                        conn_id,
960                        "embedded dyn_listen accepted a cross-process peer; embedded \
961                         multi-node forwarding uses the in-process registry. Closing \
962                         connection."
963                    );
964                    let bus = inner.bus.clone();
965                    let cancel = inner.cancel.clone();
966                    tokio::spawn(async move {
967                        let _ = sock; // drop on close
968                        let close_reason = if cancel.is_cancelled() {
969                            CloseReason::LocalClose
970                        } else {
971                            CloseReason::PeerEof
972                        };
973                        bus.send(ServerEvent::ConnectionClosed {
974                            conn_id,
975                            reason: close_reason,
976                        });
977                    });
978                }
979            }
980        }
981    }
982}
983
984async fn stats_loop(inner: Arc<ServerInner>, interval_ms: Option<i64>) {
985    let interval =
986        Duration::from_millis(u64::try_from(interval_ms.unwrap_or(1_000)).unwrap_or(1_000));
987    let mut ticker = tokio::time::interval(interval);
988    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
989    loop {
990        tokio::select! {
991            biased;
992            () = inner.cancel.cancelled() => return,
993            _ = ticker.tick() => {
994                let snap = inner.stats.snapshot();
995                *inner.snapshot_cache.lock() = snap;
996            }
997        }
998    }
999}
1000
1001async fn metrics_loop(inner: Arc<ServerInner>) {
1002    let interval = inner.metrics.flush_interval();
1003    let mut ticker = tokio::time::interval(interval.max(Duration::from_millis(50)));
1004    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
1005    loop {
1006        tokio::select! {
1007            biased;
1008            () = inner.cancel.cancelled() => return,
1009            _ = ticker.tick() => {
1010                let snap = inner.snapshot_cache.lock().clone();
1011                let _ = inner.metrics.emit(&snap).await;
1012            }
1013        }
1014    }
1015}
1016
1017async fn gossip_loop(inner: Arc<ServerInner>, interval: Duration) {
1018    let mut ticker = tokio::time::interval(interval.max(Duration::from_millis(20)));
1019    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
1020    let mut round: u64 = 0;
1021    let mut known: std::collections::HashSet<(String, u16)> = std::collections::HashSet::new();
1022    {
1023        let peers = inner.pool.peers().read();
1024        for p in peers.iter() {
1025            known.insert((p.endpoint().host().to_string(), p.endpoint().port()));
1026        }
1027    }
1028    loop {
1029        tokio::select! {
1030            biased;
1031            () = inner.cancel.cancelled() => return,
1032            _ = ticker.tick() => {
1033                let round_started = std::time::Instant::now();
1034                let round_ts = std::time::SystemTime::now();
1035                round += 1;
1036                let seeds = inner.seeds.fetch().unwrap_or_default();
1037                let mut added: u32 = 0;
1038                {
1039                    let mut peers = inner.pool.peers().write();
1040                    let now_secs = std::time::SystemTime::now()
1041                        .duration_since(std::time::UNIX_EPOCH)
1042                        .map_or(0, |d| d.as_secs());
1043                    let cfg = inner.pool.config().clone();
1044                    for seed in &seeds {
1045                        let key = (seed.host().to_string(), seed.port());
1046                        if known.contains(&key) {
1047                            continue;
1048                        }
1049                        let next_idx = u32::try_from(peers.len()).unwrap_or(u32::MAX);
1050                        let tokens: Vec<DynToken> = seed
1051                            .tokens()
1052                            .components()
1053                            .iter()
1054                            .map(|c| DynToken::from_u32(c.digits.parse::<u32>().unwrap_or(0)))
1055                            .collect();
1056                        let mut p = Peer::new(
1057                            next_idx,
1058                            PeerEndpoint::tcp(seed.host().to_string(), seed.port()),
1059                            seed.rack().to_string(),
1060                            seed.dc().to_string(),
1061                            if tokens.is_empty() {
1062                                vec![DynToken::from_u32(0)]
1063                            } else {
1064                                tokens
1065                            },
1066                            false,
1067                            seed.dc() == cfg.dc,
1068                            false,
1069                        );
1070                        p.set_state(PeerState::Normal, now_secs);
1071                        peers.push(p);
1072                        known.insert(key);
1073                        added += 1;
1074                    }
1075                }
1076                if added > 0 {
1077                    // Rebuild the topology + ring under fresh
1078                    // datacenters so the dispatcher sees the new
1079                    // peers.
1080                    rebuild_topology(&inner.pool);
1081                    let peers_now = inner.pool.peers().read();
1082                    for (idx, p) in peers_now.iter().enumerate().rev().take(added as usize) {
1083                        let _ = idx;
1084                        inner.bus.send(ServerEvent::PeerUp(p.idx()));
1085                        inner.events.publish(crate::events::ClusterEvent::PeerUp {
1086                            peer_id: p.idx(),
1087                            dc: p.dc().to_string(),
1088                            ts: round_ts,
1089                        });
1090                    }
1091                    inner.events.publish(crate::events::ClusterEvent::RingChanged {
1092                        tag: "seed-discovery".to_string(),
1093                        ts: round_ts,
1094                    });
1095                }
1096                let count = u32::try_from(inner.pool.peers().read().len()).unwrap_or(u32::MAX);
1097                inner.bus.send(ServerEvent::GossipRound { round, peers: count });
1098                inner.events.publish(crate::events::ClusterEvent::GossipRoundComplete {
1099                    duration: round_started.elapsed(),
1100                    peers_seen: count as usize,
1101                    ts: round_ts,
1102                });
1103                inner.stats.pool_incr(PoolField::StatsCount);
1104            }
1105        }
1106    }
1107}
1108
1109fn rebuild_topology(pool: &Arc<ServerPool>) {
1110    use crate::cluster::Datacenter;
1111    let peers = pool.peers().read();
1112    let mut new_dcs: Vec<Datacenter> = Vec::new();
1113    for p in peers.iter() {
1114        let idx = if let Some(i) = new_dcs.iter().position(|d| d.name() == p.dc()) {
1115            i
1116        } else {
1117            new_dcs.push(Datacenter::new(p.dc().to_string()));
1118            new_dcs.len() - 1
1119        };
1120        new_dcs[idx].upsert_rack(p.rack().to_string());
1121    }
1122    drop(peers);
1123    {
1124        let mut dcs = pool.datacenters().write();
1125        *dcs = new_dcs;
1126    }
1127    pool.rebuild_ring();
1128    pool.preselect_remote_racks();
1129}
1130
1131// Internal helpers (`shutdown_signal`, `_instant_now`) removed:
1132// they were public-but-doc(hidden) escape hatches that risked
1133// SemVer commitments. The `oneshot` channel was not used
1134// outside the (deleted) test helper, and `Instant` was kept in
1135// scope only to silence an unused-import warning. Both have
1136// been dropped.