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}