Skip to main content

dynomite/embed/
hooks.rs

1//! Pluggable hook traits exposed to embedders.
2//!
3//! The five traits in this module compose the public hook surface
4//! that an embedding program plugs into a [`crate::embed::Server`]
5//! to override the in-crate defaults: backing datastore, seeds
6//! provider, network transport listener, crypto provider, and
7//! metrics sink.
8//!
9//! Every trait is object-safe (`Box<dyn Trait>` works) and
10//! `Send + Sync` so the implementor can be shared across tokio
11//! tasks. Async methods return [`BoxFuture`] handles to keep the
12//! trait dyn-compatible without depending on the `async_trait`
13//! crate.
14//!
15//! Default implementations ship next to each trait. Re-exports
16//! at the [`crate::embed`] root mean a typical embedder writes
17//! `use dynomite::embed::SimpleSeedsProvider;` without reaching
18//! into nested submodules.
19
20use std::future::Future;
21use std::pin::Pin;
22use std::sync::Arc;
23use std::time::Duration;
24
25use parking_lot::Mutex;
26use thiserror::Error;
27
28use crate::conf::{ConfDynSeed, DataStore};
29use crate::msg::{Msg, MsgType};
30use crate::seeds::{
31    dns::DnsSeedsProvider as InnerDnsSeedsProvider,
32    florida::FloridaSeedsProvider as InnerFloridaSeedsProvider,
33    simple::SimpleSeedsProvider as InnerSimpleSeedsProvider, SeedsError, SeedsProvider as RawSeeds,
34};
35use crate::stats::{describe_stats, MetricSpec, Snapshot};
36
37/// Convenience alias for boxed futures returned by hook traits.
38///
39/// # Examples
40///
41/// ```
42/// use dynomite::embed::hooks::BoxFuture;
43/// fn _adapter() -> BoxFuture<'static, u32> {
44///     Box::pin(async { 42 })
45/// }
46/// ```
47pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
48
49// ---------- Datastore -------------------------------------------------------
50
51/// Errors produced by a [`Datastore`] implementation.
52#[derive(Debug, Error)]
53#[non_exhaustive]
54pub enum DatastoreError {
55    /// The datastore declined to handle this request type.
56    #[error("unsupported request: {0:?}")]
57    Unsupported(MsgType),
58    /// The datastore returned an internal error message.
59    #[error("datastore error: {0}")]
60    Backend(String),
61    /// I/O failure talking to the backing store.
62    #[error("io error: {0}")]
63    Io(String),
64}
65
66/// Logical wire protocol exposed by a datastore.
67///
68/// Mirrors the `data_store:` setting in the YAML config. The
69/// `Custom` variant exists for embedders fronting a private
70/// protocol; it carries no semantics inside the engine and is
71/// reported only through [`Datastore::protocol`].
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
73#[non_exhaustive]
74pub enum Protocol {
75    /// Redis RESP.
76    Redis,
77    /// Memcached text/binary.
78    Memcache,
79    /// Embedder-defined protocol.
80    Custom,
81}
82
83impl From<DataStore> for Protocol {
84    fn from(d: DataStore) -> Self {
85        match d {
86            DataStore::Valkey => Protocol::Redis,
87            DataStore::Memcache => Protocol::Memcache,
88            DataStore::Dyniak => Protocol::Custom,
89        }
90    }
91}
92
93/// Backing datastore the engine forwards routed requests to.
94///
95/// Implementations are kept stateless on the trait surface; live
96/// connection state lives behind [`Datastore::dispatch`] in the
97/// implementor's chosen pool.
98///
99/// # Examples
100///
101/// ```
102/// use dynomite::embed::hooks::{Datastore, MemoryDatastore, Protocol};
103/// use dynomite::msg::{Msg, MsgType};
104/// # tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(async {
105/// let ds = MemoryDatastore::new();
106/// assert_eq!(ds.protocol(), Protocol::Custom);
107/// let req = Msg::new(1, MsgType::ReqRedisGet, true);
108/// let _rsp = ds.dispatch(req).await.unwrap();
109/// assert_eq!(ds.dispatch_count(), 1);
110/// # });
111/// ```
112pub trait Datastore: Send + Sync {
113    /// Return the wire protocol the datastore speaks.
114    fn protocol(&self) -> Protocol;
115
116    /// Predicate used by the dispatcher to short-circuit
117    /// commands the backend cannot serve.
118    fn supports(&self, _cmd: MsgType) -> bool {
119        true
120    }
121
122    /// Forward a routed request and return the response message.
123    ///
124    /// The returned future is `'static`; the caller may move it
125    /// across tokio tasks freely.
126    fn dispatch(&self, req: Msg) -> BoxFuture<'_, Result<Msg, DatastoreError>>;
127
128    /// Optional downcast hook for transports that need a concrete
129    /// backend capability the trait surface does not express.
130    ///
131    /// The default returns `None`, so an existing impl that has no
132    /// extra capability does not have to override anything. A
133    /// backend that wants to expose itself for downcasting (for
134    /// example the Riak multi-key transaction store, which a
135    /// transport probes for via [`std::any::Any`] `downcast_ref`)
136    /// returns `Some(self)`.
137    fn as_any(&self) -> Option<&(dyn std::any::Any + 'static)> {
138        None
139    }
140
141    /// Stream the names of every bucket the datastore is aware of,
142    /// one [`bytes::Bytes`] per bucket.
143    ///
144    /// The default implementation returns a one-shot stream that
145    /// yields a single [`DatastoreError::Unsupported`] item, so an
146    /// existing impl that does not enumerate buckets does not have
147    /// to be modified to compile against the new trait surface.
148    /// Implementations that can enumerate override this method.
149    ///
150    /// The returned stream is `'static` and `Send`; transports that
151    /// stream chunks of bucket names to clients (PBC frames, HTTP
152    /// chunked bodies, ...) consume it from a tokio task.
153    fn list_buckets_stream(&self) -> DatastoreByteStream {
154        Box::pin(unsupported_byte_stream())
155    }
156
157    /// Stream every key in `bucket`, one [`bytes::Bytes`] per key.
158    ///
159    /// Same default as [`Datastore::list_buckets_stream`]: a
160    /// single-item stream carrying [`DatastoreError::Unsupported`].
161    fn list_keys_stream(&self, _bucket: &[u8]) -> DatastoreByteStream {
162        Box::pin(unsupported_byte_stream())
163    }
164
165    /// Read the object stored under `(bucket, key)` against the
166    /// Riak K/V layer.
167    ///
168    /// Returns `Ok(None)` when no object exists. The default
169    /// implementation reports the operation as unsupported so
170    /// existing `Datastore` impls that do not speak Riak
171    /// continue to compile against the new trait surface; the
172    /// PBC server treats the error as a routing failure and
173    /// emits an `RpbErrorResp`.
174    fn riak_get<'a>(
175        &'a self,
176        _bucket: &'a [u8],
177        _key: &'a [u8],
178    ) -> BoxFuture<'a, Result<Option<Vec<u8>>, DatastoreError>> {
179        Box::pin(async move { Err(DatastoreError::Unsupported(MsgType::Unknown)) })
180    }
181
182    /// Store `value` under `(bucket, key)`. `indexes` carries
183    /// `(index_name, encoded_value)` pairs to associate with the
184    /// object on the 2i layer.
185    ///
186    /// Default: unsupported, see [`Datastore::riak_get`].
187    fn riak_put<'a>(
188        &'a self,
189        _bucket: &'a [u8],
190        _key: &'a [u8],
191        _value: &'a [u8],
192        _indexes: &'a [(Vec<u8>, Vec<u8>)],
193    ) -> BoxFuture<'a, Result<(), DatastoreError>> {
194        Box::pin(async move { Err(DatastoreError::Unsupported(MsgType::Unknown)) })
195    }
196
197    /// Delete the object stored under `(bucket, key)`. Returns
198    /// `true` when an object was removed, `false` when none
199    /// existed.
200    ///
201    /// Default: unsupported, see [`Datastore::riak_get`].
202    fn riak_delete<'a>(
203        &'a self,
204        _bucket: &'a [u8],
205        _key: &'a [u8],
206    ) -> BoxFuture<'a, Result<bool, DatastoreError>> {
207        Box::pin(async move { Err(DatastoreError::Unsupported(MsgType::Unknown)) })
208    }
209
210    /// Equality query against the 2i layer.
211    ///
212    /// Returns the object keys whose `index_name` value equals
213    /// `value`, ordered by the underlying storage's natural key
214    /// order (typically lexicographic).
215    ///
216    /// Default: unsupported, see [`Datastore::riak_get`].
217    fn riak_index_eq<'a>(
218        &'a self,
219        _bucket: &'a [u8],
220        _index_name: &'a [u8],
221        _value: &'a [u8],
222    ) -> BoxFuture<'a, Result<Vec<Vec<u8>>, DatastoreError>> {
223        Box::pin(async move { Err(DatastoreError::Unsupported(MsgType::Unknown)) })
224    }
225
226    /// Range query against the 2i layer. `min` and `max` are
227    /// inclusive bounds in the same encoding the index uses
228    /// internally.
229    ///
230    /// Default: unsupported, see [`Datastore::riak_get`].
231    fn riak_index_range<'a>(
232        &'a self,
233        _bucket: &'a [u8],
234        _index_name: &'a [u8],
235        _min: &'a [u8],
236        _max: &'a [u8],
237    ) -> BoxFuture<'a, Result<Vec<Vec<u8>>, DatastoreError>> {
238        Box::pin(async move { Err(DatastoreError::Unsupported(MsgType::Unknown)) })
239    }
240}
241
242/// In-memory datastore used by examples and integration tests.
243///
244/// Stores the "command -> response" pairing for the simplest
245/// commands without speaking the real protocol. Production
246/// embedders use [`RedisDatastore`] or [`MemcacheDatastore`].
247#[derive(Debug, Default, Clone)]
248pub struct MemoryDatastore {
249    inner: Arc<Mutex<MemoryStore>>,
250}
251
252#[derive(Debug, Default)]
253struct MemoryStore {
254    calls: u64,
255}
256
257impl MemoryDatastore {
258    /// Build a fresh empty store.
259    ///
260    /// # Examples
261    ///
262    /// ```
263    /// use dynomite::embed::hooks::MemoryDatastore;
264    /// let ds = MemoryDatastore::new();
265    /// assert_eq!(ds.dispatch_count(), 0);
266    /// ```
267    #[must_use]
268    pub fn new() -> Self {
269        Self::default()
270    }
271
272    /// Number of times [`Datastore::dispatch`] has been invoked.
273    ///
274    /// # Examples
275    ///
276    /// ```
277    /// use dynomite::embed::hooks::MemoryDatastore;
278    /// let ds = MemoryDatastore::new();
279    /// assert_eq!(ds.dispatch_count(), 0);
280    /// ```
281    #[must_use]
282    pub fn dispatch_count(&self) -> u64 {
283        self.inner.lock().calls
284    }
285}
286
287impl Datastore for MemoryDatastore {
288    fn protocol(&self) -> Protocol {
289        Protocol::Custom
290    }
291
292    fn dispatch(&self, req: Msg) -> BoxFuture<'_, Result<Msg, DatastoreError>> {
293        let inner = self.inner.clone();
294        Box::pin(async move {
295            inner.lock().calls += 1;
296            let mut rsp = Msg::new(req.id(), MsgType::Unknown, false);
297            rsp.set_parent_id(req.id());
298            Ok(rsp)
299        })
300    }
301
302    fn list_buckets_stream(&self) -> DatastoreByteStream {
303        let snapshot = self.list_buckets_snapshot();
304        Box::pin(VecByteStream {
305            items: snapshot.into_iter(),
306        })
307    }
308
309    fn list_keys_stream(&self, bucket: &[u8]) -> DatastoreByteStream {
310        let snapshot = self.list_keys_snapshot(bucket);
311        Box::pin(VecByteStream {
312            items: snapshot.into_iter(),
313        })
314    }
315}
316
317/// Default RESP-fronting datastore.
318///
319/// A thin marker around the supplied
320/// connection target; the actual wire protocol bridge lives in
321/// the dispatcher path of [`crate::cluster::dispatch`]. The
322/// default impl satisfies the [`Datastore`] contract for the
323/// embed surface so an embedder can construct a builder without
324/// wiring a custom backend.
325#[derive(Debug, Clone)]
326pub struct RedisDatastore {
327    target: String,
328}
329
330impl RedisDatastore {
331    /// Build a new Redis-fronting datastore.
332    ///
333    /// `target` is informational and is reported back through
334    /// [`RedisDatastore::target`].
335    ///
336    /// # Examples
337    ///
338    /// ```
339    /// use dynomite::embed::hooks::RedisDatastore;
340    /// let r = RedisDatastore::new("127.0.0.1:6379");
341    /// assert_eq!(r.target(), "127.0.0.1:6379");
342    /// ```
343    pub fn new(target: impl Into<String>) -> Self {
344        Self {
345            target: target.into(),
346        }
347    }
348
349    /// Return the configured target string.
350    #[must_use]
351    pub fn target(&self) -> &str {
352        &self.target
353    }
354}
355
356impl Datastore for RedisDatastore {
357    fn protocol(&self) -> Protocol {
358        Protocol::Redis
359    }
360
361    fn dispatch(&self, req: Msg) -> BoxFuture<'_, Result<Msg, DatastoreError>> {
362        Box::pin(async move {
363            let mut rsp = Msg::new(req.id(), MsgType::RspRedisStatus, false);
364            rsp.set_parent_id(req.id());
365            Ok(rsp)
366        })
367    }
368}
369
370/// Default Memcache-fronting datastore.
371///
372/// Mirrors [`RedisDatastore`] for the Memcache wire protocol.
373#[derive(Debug, Clone)]
374pub struct MemcacheDatastore {
375    target: String,
376}
377
378impl MemcacheDatastore {
379    /// Build a new Memcache-fronting datastore.
380    ///
381    /// # Examples
382    ///
383    /// ```
384    /// use dynomite::embed::hooks::MemcacheDatastore;
385    /// let m = MemcacheDatastore::new("127.0.0.1:11211");
386    /// assert_eq!(m.target(), "127.0.0.1:11211");
387    /// ```
388    pub fn new(target: impl Into<String>) -> Self {
389        Self {
390            target: target.into(),
391        }
392    }
393
394    /// Return the configured target string.
395    #[must_use]
396    pub fn target(&self) -> &str {
397        &self.target
398    }
399}
400
401impl Datastore for MemcacheDatastore {
402    fn protocol(&self) -> Protocol {
403        Protocol::Memcache
404    }
405
406    fn dispatch(&self, req: Msg) -> BoxFuture<'_, Result<Msg, DatastoreError>> {
407        Box::pin(async move {
408            let mut rsp = Msg::new(req.id(), MsgType::RspMcEnd, false);
409            rsp.set_parent_id(req.id());
410            Ok(rsp)
411        })
412    }
413}
414
415// ---------- SeedsProvider ---------------------------------------------------
416
417/// Pluggable seeds provider.
418///
419/// The trait is the embed-API mirror of
420/// [`crate::seeds::SeedsProvider`]; the in-crate providers are
421/// re-exported below as default implementations.
422///
423/// # Examples
424///
425/// ```
426/// use dynomite::embed::hooks::{SeedsProvider, SimpleSeedsProvider};
427/// use dynomite::conf::ConfDynSeed;
428/// let sp = SimpleSeedsProvider::new(vec![ConfDynSeed::parse("h:1:r:d:1").unwrap()]);
429/// assert_eq!(sp.fetch().unwrap().len(), 1);
430/// ```
431pub trait SeedsProvider: Send + Sync {
432    /// Return the current list of seeds.
433    fn fetch(&self) -> Result<Vec<ConfDynSeed>, SeedsError>;
434
435    /// Refresh interval used by the gossip task between calls to
436    /// [`SeedsProvider::fetch`].
437    fn refresh_interval(&self) -> Duration {
438        Duration::from_secs(30)
439    }
440}
441
442// LegacySeedsAdapter removed: the type would have leaked the
443// in-crate `crate::seeds::SeedsProvider` trait onto the public
444// API via its generic bound. Embedders that need to lift an
445// existing `SeedsProvider` impl wrap it in
446// `Box<dyn SeedsProvider>` directly and pass it to
447// [`crate::embed::ServerBuilder::seeds_provider`].
448
449/// Re-exported [`crate::seeds::simple::SimpleSeedsProvider`] with
450/// the embed [`SeedsProvider`] trait already implemented.
451#[derive(Debug, Clone, Default)]
452pub struct SimpleSeedsProvider {
453    inner: InnerSimpleSeedsProvider,
454}
455
456impl SimpleSeedsProvider {
457    /// Build a new in-memory provider.
458    ///
459    /// # Examples
460    ///
461    /// ```
462    /// use dynomite::embed::hooks::{SeedsProvider, SimpleSeedsProvider};
463    /// let p = SimpleSeedsProvider::new(Vec::new());
464    /// assert_eq!(p.fetch().unwrap().len(), 0);
465    /// ```
466    #[must_use]
467    pub fn new(seeds: Vec<ConfDynSeed>) -> Self {
468        Self {
469            inner: InnerSimpleSeedsProvider::new(seeds),
470        }
471    }
472}
473
474impl SeedsProvider for SimpleSeedsProvider {
475    fn fetch(&self) -> Result<Vec<ConfDynSeed>, SeedsError> {
476        self.inner.get_seeds()
477    }
478}
479
480/// DNS-resolving seeds provider, re-exported under the embed
481/// trait.
482pub type DnsSeedsProvider = InnerDnsSeedsProvider;
483
484impl SeedsProvider for DnsSeedsProvider {
485    fn fetch(&self) -> Result<Vec<ConfDynSeed>, SeedsError> {
486        self.get_seeds()
487    }
488}
489
490/// Florida HTTP seeds provider, re-exported under the embed
491/// trait.
492pub type FloridaSeedsProvider = InnerFloridaSeedsProvider;
493
494impl SeedsProvider for FloridaSeedsProvider {
495    fn fetch(&self) -> Result<Vec<ConfDynSeed>, SeedsError> {
496        self.get_seeds()
497    }
498}
499
500// ---------- CryptoProvider --------------------------------------------------
501
502/// Errors produced by a [`CryptoProvider`].
503#[derive(Debug, Error)]
504#[non_exhaustive]
505pub enum CryptoProviderError {
506    /// Encryption failed.
507    #[error("encryption failed: {0}")]
508    Encrypt(String),
509    /// Decryption failed.
510    #[error("decryption failed: {0}")]
511    Decrypt(String),
512    /// The provider was misconfigured (key length, padding, ...).
513    #[error("misconfiguration: {0}")]
514    Misconfigured(String),
515}
516
517/// Pluggable AES + RSA provider for the DNODE peer protocol.
518///
519/// The default in-crate provider [`RustCryptoProvider`] wraps
520/// [`crate::crypto::Crypto`]. HSM / KMS integrations implement
521/// the trait against their hardware bridge.
522///
523/// # Examples
524///
525/// ```no_run
526/// use dynomite::embed::hooks::{CryptoProvider, RustCryptoProvider};
527/// use dynomite::crypto::Crypto;
528/// // Construct the underlying Crypto from a PEM file at runtime.
529/// let crypto = Crypto::from_pem("/etc/dynomite/dynomite.pem").unwrap();
530/// let provider = RustCryptoProvider::new(crypto);
531/// assert!(provider.rsa_size() > 0);
532/// ```
533pub trait CryptoProvider: Send + Sync {
534    /// Modulus length in bytes for the configured RSA key.
535    fn rsa_size(&self) -> usize;
536
537    /// Borrow the AES key buffer used by the DNODE handshake.
538    fn aes_key(&self) -> [u8; crate::crypto::AES_KEYLEN];
539
540    /// AES-encrypt `plaintext` under the provider's key.
541    fn aes_encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoProviderError>;
542
543    /// AES-decrypt `ciphertext` under the provider's key.
544    fn aes_decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoProviderError>;
545
546    /// RSA-encrypt `plaintext` under the provider's public key.
547    fn rsa_encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoProviderError>;
548
549    /// RSA-decrypt `ciphertext` under the provider's private key.
550    fn rsa_decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoProviderError>;
551}
552
553/// Default crypto provider built on the in-crate
554/// [`crate::crypto::Crypto`] (RustCrypto-backed AES + RSA).
555///
556/// Despite the historical "Openssl" name in the design doc, the
557/// implementation uses the workspace's RustCrypto stack. Recorded
558/// as a Deviation in `docs/parity.md`.
559#[derive(Debug)]
560pub struct RustCryptoProvider {
561    crypto: Arc<crate::crypto::Crypto>,
562}
563
564impl RustCryptoProvider {
565    /// Wrap an existing [`crate::crypto::Crypto`].
566    ///
567    /// # Examples
568    ///
569    /// ```no_run
570    /// use dynomite::embed::hooks::{CryptoProvider, RustCryptoProvider};
571    /// use dynomite::crypto::Crypto;
572    /// // Production embedders load the key from disk:
573    /// let crypto = Crypto::from_pem("/etc/dynomite/dynomite.pem").unwrap();
574    /// let p = RustCryptoProvider::new(crypto);
575    /// assert!(p.rsa_size() >= 128);
576    /// ```
577    #[must_use]
578    pub fn new(crypto: crate::crypto::Crypto) -> Self {
579        Self {
580            crypto: Arc::new(crypto),
581        }
582    }
583
584    /// Construct from an [`Arc`]-shared crypto bundle.
585    #[must_use]
586    pub fn from_arc(crypto: Arc<crate::crypto::Crypto>) -> Self {
587        Self { crypto }
588    }
589}
590
591impl CryptoProvider for RustCryptoProvider {
592    fn rsa_size(&self) -> usize {
593        self.crypto.rsa_size()
594    }
595
596    fn aes_key(&self) -> [u8; crate::crypto::AES_KEYLEN] {
597        *self.crypto.aes_key()
598    }
599
600    fn aes_encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoProviderError> {
601        crate::crypto::Crypto::aes_encrypt(plaintext, self.crypto.aes_key())
602            .map_err(|e| CryptoProviderError::Encrypt(e.to_string()))
603    }
604
605    fn aes_decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoProviderError> {
606        crate::crypto::Crypto::aes_decrypt(ciphertext, self.crypto.aes_key())
607            .map_err(|e| CryptoProviderError::Decrypt(e.to_string()))
608    }
609
610    fn rsa_encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoProviderError> {
611        self.crypto
612            .rsa_encrypt(plaintext)
613            .map_err(|e| CryptoProviderError::Encrypt(e.to_string()))
614    }
615
616    fn rsa_decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoProviderError> {
617        self.crypto
618            .rsa_decrypt(ciphertext)
619            .map_err(|e| CryptoProviderError::Decrypt(e.to_string()))
620    }
621}
622
623// ---------- MetricsSink -----------------------------------------------------
624
625/// Errors produced by a [`MetricsSink`].
626#[derive(Debug, Error)]
627#[non_exhaustive]
628pub enum MetricsError {
629    /// The sink could not flush the snapshot.
630    #[error("metrics flush failed: {0}")]
631    Flush(String),
632}
633
634/// Pluggable metrics exporter.
635///
636/// The default sink ([`LoggingMetricsSink`]) emits one `tracing`
637/// event per flush. A `PrometheusMetricsSink` is intentionally
638/// not shipped by default to avoid adding a dependency to the
639/// workspace; see the embedding cookbook for the recommended
640/// adapter shape. Recorded as a Deviation in `docs/parity.md`.
641///
642/// # Examples
643///
644/// ```
645/// use dynomite::embed::hooks::{LoggingMetricsSink, MetricsSink};
646/// use dynomite::stats::Snapshot;
647/// # tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(async {
648/// let sink = LoggingMetricsSink::new("test");
649/// sink.emit(&Snapshot::default()).await.unwrap();
650/// # });
651/// ```
652pub trait MetricsSink: Send + Sync {
653    /// Push a stats snapshot to the sink.
654    fn emit<'a>(&'a self, snapshot: &'a Snapshot) -> BoxFuture<'a, Result<(), MetricsError>>;
655
656    /// Flush interval requested by the sink.
657    fn flush_interval(&self) -> Duration {
658        Duration::from_secs(10)
659    }
660
661    /// Optional manifest of every metric the sink expects to
662    /// receive.
663    fn manifest(&self) -> Vec<MetricSpec> {
664        Vec::new()
665    }
666}
667
668/// Default metrics sink: emits one `tracing::info` event per
669/// flush. Cheap, dependency-free, and useful in development.
670#[derive(Debug, Clone)]
671pub struct LoggingMetricsSink {
672    name: String,
673    counter: Arc<Mutex<u64>>,
674}
675
676impl LoggingMetricsSink {
677    /// Build a sink with a human-readable name.
678    ///
679    /// # Examples
680    ///
681    /// ```
682    /// use dynomite::embed::hooks::LoggingMetricsSink;
683    /// let s = LoggingMetricsSink::new("dyn_o_mite");
684    /// assert_eq!(s.flush_count(), 0);
685    /// ```
686    #[must_use]
687    pub fn new(name: impl Into<String>) -> Self {
688        Self {
689            name: name.into(),
690            counter: Arc::new(Mutex::new(0)),
691        }
692    }
693
694    /// Number of flushes the sink has observed.
695    #[must_use]
696    pub fn flush_count(&self) -> u64 {
697        *self.counter.lock()
698    }
699}
700
701impl MetricsSink for LoggingMetricsSink {
702    fn emit<'a>(&'a self, snapshot: &'a Snapshot) -> BoxFuture<'a, Result<(), MetricsError>> {
703        let counter = self.counter.clone();
704        let name = self.name.clone();
705        let pool_name = snapshot.pool.name.clone();
706        Box::pin(async move {
707            *counter.lock() += 1;
708            tracing::info!(sink = %name, pool = %pool_name, "metrics flush");
709            Ok(())
710        })
711    }
712
713    fn manifest(&self) -> Vec<MetricSpec> {
714        // Defer to the in-crate descriptor table; the JSON form
715        // already contains every metric the engine emits.
716        let json = describe_stats();
717        // The descriptor table is consumed downstream as opaque
718        // text. Returning an empty vec keeps the trait signature
719        // honest while logging the manifest size for diagnostics.
720        tracing::debug!(bytes = json.len(), "metrics manifest");
721        Vec::new()
722    }
723}
724
725// ---------- Streaming Datastore extensions ---------------------------------
726//
727// The streaming list path lets transports emit the bucket / key
728// catalogue in chunks instead of buffering it. A `Datastore` impl
729// produces one `Bytes` per entry; the transport (PBC server, HTTP
730// gateway) buffers an implementation-chosen number of entries per
731// outbound frame.
732
733use bytes::Bytes;
734use std::collections::{BTreeMap, BTreeSet};
735use std::sync::OnceLock;
736use std::task::{Context, Poll};
737
738/// Type alias for the byte stream returned by
739/// [`Datastore::list_buckets_stream`] and
740/// [`Datastore::list_keys_stream`].
741///
742/// Each `Bytes` item is one bucket name or key. The stream may
743/// surface a [`DatastoreError`] mid-iteration; transports translate
744/// that into a wire-level error response and stop emitting frames.
745pub type DatastoreByteStream =
746    Pin<Box<dyn futures_core::Stream<Item = Result<Bytes, DatastoreError>> + Send>>;
747
748/// Build a one-shot stream that yields a single
749/// [`DatastoreError::Unsupported`] item. The default body for
750/// [`Datastore::list_buckets_stream`] /
751/// [`Datastore::list_keys_stream`] so existing impls keep working
752/// without modification.
753fn unsupported_byte_stream(
754) -> impl futures_core::Stream<Item = Result<Bytes, DatastoreError>> + Send {
755    UnsupportedListStream { emitted: false }
756}
757
758struct UnsupportedListStream {
759    emitted: bool,
760}
761
762impl futures_core::Stream for UnsupportedListStream {
763    type Item = Result<Bytes, DatastoreError>;
764
765    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
766        if self.emitted {
767            return Poll::Ready(None);
768        }
769        self.emitted = true;
770        Poll::Ready(Some(Err(DatastoreError::Unsupported(MsgType::Unknown))))
771    }
772}
773
774/// `futures_core::Stream` that drains an owned `Vec<Bytes>` one
775/// item at a time. Used by [`MemoryDatastore`] to back its
776/// streaming list overrides.
777struct VecByteStream {
778    items: std::vec::IntoIter<Bytes>,
779}
780
781impl futures_core::Stream for VecByteStream {
782    type Item = Result<Bytes, DatastoreError>;
783
784    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
785        match self.items.next() {
786            Some(b) => Poll::Ready(Some(Ok(b))),
787            None => Poll::Ready(None),
788        }
789    }
790}
791
792// ---------- MemoryDatastore listing index ----------------------------------
793//
794// `MemoryDatastore`'s public field layout (a single `Arc<Mutex<...>>`
795// holding the dispatch counter) was committed before the streaming
796// list path existed. To keep the published surface unchanged, the
797// per-instance bucket/key index is held in a process-wide registry
798// keyed by the `Arc<Mutex<MemoryStore>>` pointer identity. Cloning
799// a `MemoryDatastore` shares its `inner` `Arc`, so clones see the
800// same listing -- mirroring the existing dispatch-count behaviour.
801
802#[derive(Debug, Default)]
803struct MemoryListing {
804    buckets: BTreeMap<Vec<u8>, BTreeSet<Vec<u8>>>,
805}
806
807#[derive(Debug, Default, Clone)]
808struct ListingHandle {
809    inner: Arc<Mutex<MemoryListing>>,
810}
811
812fn listing_for(ds: &MemoryDatastore) -> ListingHandle {
813    static REGISTRY: OnceLock<Mutex<Vec<(usize, ListingHandle)>>> = OnceLock::new();
814    let registry = REGISTRY.get_or_init(|| Mutex::new(Vec::new()));
815    let id = Arc::as_ptr(&ds.inner) as usize;
816    let mut g = registry.lock();
817    if let Some((_, h)) = g.iter().find(|(k, _)| *k == id) {
818        return h.clone();
819    }
820    let h = ListingHandle::default();
821    g.push((id, h.clone()));
822    h
823}
824
825impl MemoryDatastore {
826    /// Insert `(bucket, key)` into the in-memory listing index.
827    ///
828    /// Idempotent: inserting the same `(bucket, key)` pair twice is
829    /// a no-op. Tests use this helper to seed the streaming list
830    /// path without speaking the real Riak protocol.
831    ///
832    /// # Examples
833    ///
834    /// ```
835    /// use dynomite::embed::hooks::MemoryDatastore;
836    /// let ds = MemoryDatastore::new();
837    /// ds.insert(b"users", b"alice");
838    /// ds.insert(b"users", b"bob");
839    /// assert_eq!(ds.list_buckets_snapshot().len(), 1);
840    /// assert_eq!(ds.list_keys_snapshot(b"users").len(), 2);
841    /// ```
842    pub fn insert(&self, bucket: &[u8], key: &[u8]) {
843        let h = listing_for(self);
844        let mut g = h.inner.lock();
845        g.buckets
846            .entry(bucket.to_vec())
847            .or_default()
848            .insert(key.to_vec());
849    }
850
851    /// Snapshot of the bucket name set, sorted lexicographically.
852    ///
853    /// # Examples
854    ///
855    /// ```
856    /// use dynomite::embed::hooks::MemoryDatastore;
857    /// let ds = MemoryDatastore::new();
858    /// assert!(ds.list_buckets_snapshot().is_empty());
859    /// ```
860    #[must_use]
861    pub fn list_buckets_snapshot(&self) -> Vec<Bytes> {
862        let h = listing_for(self);
863        let g = h.inner.lock();
864        g.buckets
865            .keys()
866            .map(|b| Bytes::copy_from_slice(b))
867            .collect()
868    }
869
870    /// Snapshot of the keys in `bucket`, sorted lexicographically.
871    /// Returns an empty vector when the bucket is unknown.
872    ///
873    /// # Examples
874    ///
875    /// ```
876    /// use dynomite::embed::hooks::MemoryDatastore;
877    /// let ds = MemoryDatastore::new();
878    /// assert!(ds.list_keys_snapshot(b"missing").is_empty());
879    /// ```
880    #[must_use]
881    pub fn list_keys_snapshot(&self, bucket: &[u8]) -> Vec<Bytes> {
882        let h = listing_for(self);
883        let g = h.inner.lock();
884        g.buckets
885            .get(bucket)
886            .map(|s| s.iter().map(|k| Bytes::copy_from_slice(k)).collect())
887            .unwrap_or_default()
888    }
889}