1mod file;
4mod memory;
5pub use file::{FileNodeStore, FileNodeStoreError};
6pub use memory::{MemStore, MemStoreError};
7
8use std::collections::{hash_map::Entry, HashMap};
9use std::future::Future;
10use std::sync::Arc;
11
12use super::cid::Cid;
13use super::engine::ready::{ready_only, ReadyOnly};
14use super::manifest::{
15 AsyncManifestStore, AsyncManifestStoreScan, ManifestStore, ManifestStoreScan, ManifestUpdate,
16 NamedRootManifest, RootManifest,
17};
18
19pub(crate) struct OrderedBatchReadPlan<'a> {
20 unique_keys: Vec<&'a [u8]>,
21 positions: Option<Vec<usize>>,
22}
23
24impl<'a> OrderedBatchReadPlan<'a> {
25 pub(crate) fn new(keys: &[&'a [u8]]) -> Self {
26 if keys.len() < 2 {
27 return Self {
28 unique_keys: keys.to_vec(),
29 positions: None,
30 };
31 }
32
33 let mut unique_indexes = HashMap::with_capacity(keys.len());
34 let mut unique_keys = Vec::with_capacity(keys.len());
35 let mut positions: Option<Vec<usize>> = None;
36
37 for key in keys {
38 match unique_indexes.entry(*key) {
39 Entry::Occupied(entry) => {
40 let positions =
41 positions.get_or_insert_with(|| (0..unique_keys.len()).collect());
42 positions.push(*entry.get());
43 }
44 Entry::Vacant(entry) => {
45 let unique_idx = unique_keys.len();
46 unique_keys.push(*key);
47 if let Some(positions) = positions.as_mut() {
48 positions.push(unique_idx);
49 }
50 entry.insert(unique_idx);
51 }
52 }
53 }
54
55 Self {
56 unique_keys,
57 positions,
58 }
59 }
60
61 pub(crate) fn unique_keys(&self) -> &[&'a [u8]] {
62 &self.unique_keys
63 }
64
65 #[cfg(test)]
66 pub(crate) fn is_identity(&self) -> bool {
67 self.positions.is_none()
68 }
69
70 #[cfg(test)]
71 pub(crate) fn expand<T: Clone>(&self, unique_values: &[Option<T>]) -> Vec<Option<T>> {
72 debug_assert_eq!(self.unique_keys.len(), unique_values.len());
73 match &self.positions {
74 Some(positions) => positions
75 .iter()
76 .map(|&unique_idx| unique_values[unique_idx].clone())
77 .collect(),
78 None => unique_values.to_vec(),
79 }
80 }
81
82 pub(crate) fn expand_owned<T: Clone>(&self, unique_values: Vec<Option<T>>) -> Vec<Option<T>> {
83 debug_assert_eq!(self.unique_keys.len(), unique_values.len());
84 match &self.positions {
85 Some(positions) => positions
86 .iter()
87 .map(|&unique_idx| unique_values[unique_idx].clone())
88 .collect(),
89 None => unique_values,
90 }
91 }
92}
93
94#[derive(Debug, Clone)]
96pub enum BatchOp<'a> {
97 Upsert { key: &'a [u8], value: &'a [u8] },
99 Delete { key: &'a [u8] },
101}
102
103#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
110#[non_exhaustive]
111pub enum PublicationOrigin {
112 #[default]
114 General,
115 PointUpsert,
117 PointDelete,
119 BatchMutation,
121 TreeBuild,
123 Merge,
125 RangeDelete,
127 Replication,
129 Maintenance,
131}
132
133#[derive(Clone, Copy, Debug, Eq, PartialEq)]
138pub struct NodePublicationHint<'a> {
139 namespace: &'a [u8],
140 key: &'a [u8],
141 value: &'a [u8],
142}
143
144impl<'a> NodePublicationHint<'a> {
145 #[inline(always)]
147 pub const fn new(namespace: &'a [u8], key: &'a [u8], value: &'a [u8]) -> Self {
148 Self {
149 namespace,
150 key,
151 value,
152 }
153 }
154
155 #[inline(always)]
157 pub const fn namespace(&self) -> &'a [u8] {
158 self.namespace
159 }
160
161 #[inline(always)]
163 pub const fn key(&self) -> &'a [u8] {
164 self.key
165 }
166
167 #[inline(always)]
169 pub const fn value(&self) -> &'a [u8] {
170 self.value
171 }
172}
173
174#[derive(Clone, Copy, Debug, Eq, PartialEq)]
181pub struct NodePublication<'a> {
182 entries: &'a [(&'a [u8], &'a [u8])],
183 hint: Option<NodePublicationHint<'a>>,
184 origin: PublicationOrigin,
185}
186
187impl<'a> NodePublication<'a> {
188 #[inline(always)]
190 pub const fn new(entries: &'a [(&'a [u8], &'a [u8])], origin: PublicationOrigin) -> Self {
191 Self {
192 entries,
193 hint: None,
194 origin,
195 }
196 }
197
198 #[inline(always)]
200 pub const fn with_hint(
201 entries: &'a [(&'a [u8], &'a [u8])],
202 hint: NodePublicationHint<'a>,
203 origin: PublicationOrigin,
204 ) -> Self {
205 Self {
206 entries,
207 hint: Some(hint),
208 origin,
209 }
210 }
211
212 #[inline(always)]
214 pub const fn entries(&self) -> &'a [(&'a [u8], &'a [u8])] {
215 self.entries
216 }
217
218 #[inline(always)]
220 pub const fn hint(&self) -> Option<NodePublicationHint<'a>> {
221 self.hint
222 }
223
224 #[inline(always)]
226 pub const fn origin(&self) -> PublicationOrigin {
227 self.origin
228 }
229}
230
231pub type SharedReadBatch = Vec<Option<Arc<[u8]>>>;
233
234pub trait Store: Send + Sync {
239 type Error: std::error::Error + Send + Sync + 'static;
241
242 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error>;
247
248 fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
252 self.get(key)
253 .map(|value| value.map(|bytes| Arc::from(bytes.into_boxed_slice())))
254 }
255
256 fn batch_get_shared_ordered_unique(
258 &self,
259 keys: &[&[u8]],
260 ) -> Result<SharedReadBatch, Self::Error> {
261 self.batch_get_ordered_unique(keys).map(|values| {
262 values
263 .into_iter()
264 .map(|value| value.map(|bytes| Arc::from(bytes.into_boxed_slice())))
265 .collect()
266 })
267 }
268
269 fn has_native_shared_reads(&self) -> bool {
271 false
272 }
273
274 fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error>;
278
279 fn delete(&self, key: &[u8]) -> Result<(), Self::Error>;
283
284 fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error>;
289
290 fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
298 let plan = OrderedBatchReadPlan::new(keys);
299 let mut results = HashMap::with_capacity(plan.unique_keys().len());
300 for key in plan.unique_keys() {
301 if let Some(value) = self.get(key)? {
302 results.insert(key.to_vec(), value);
303 }
304 }
305 Ok(results)
306 }
307
308 fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
327 if keys.len() < 2 {
328 return keys.iter().map(|key| self.get(key)).collect();
329 }
330
331 let plan = OrderedBatchReadPlan::new(keys);
332 let mut unique_values = Vec::with_capacity(plan.unique_keys().len());
333 for key in plan.unique_keys() {
334 unique_values.push(self.get(key)?);
335 }
336 Ok(plan.expand_owned(unique_values))
337 }
338
339 fn batch_get_ordered_unique(
346 &self,
347 keys: &[&[u8]],
348 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
349 if keys.is_empty() {
350 return Ok(Vec::new());
351 }
352
353 if !self.prefers_batch_reads() {
354 return keys.iter().map(|key| self.get(key)).collect();
355 }
356
357 self.batch_get_ordered(keys)
358 }
359
360 fn prefers_batch_reads(&self) -> bool {
366 false
367 }
368
369 fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
376 let ops: Vec<BatchOp> = entries
377 .iter()
378 .map(|(k, v)| BatchOp::Upsert { key: k, value: v })
379 .collect();
380 self.batch(&ops)
381 }
382
383 fn supports_hints(&self) -> bool {
385 false
386 }
387
388 fn prefers_rightmost_path_hints(&self) -> bool {
394 false
395 }
396
397 fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
403 let _ = (namespace, key);
404 Ok(None)
405 }
406
407 fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
411 let _ = (namespace, key, value);
412 Ok(())
413 }
414
415 fn batch_put_with_hint(
417 &self,
418 entries: &[(&[u8], &[u8])],
419 namespace: &[u8],
420 key: &[u8],
421 value: &[u8],
422 ) -> Result<(), Self::Error> {
423 self.batch_put(entries)?;
424 self.put_hint(namespace, key, value)
425 }
426
427 #[inline(always)]
433 fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
434 match publication.hint.as_ref() {
435 Some(hint) => self.batch_put_with_hint(
436 publication.entries,
437 hint.namespace(),
438 hint.key(),
439 hint.value(),
440 ),
441 None => self.batch_put(publication.entries),
442 }
443 }
444}
445
446pub trait NodeStoreScan: Send + Sync {
452 type Error: std::error::Error + Send + Sync + 'static;
454
455 fn list_node_cids(&self) -> Result<Vec<Cid>, Self::Error>;
461}
462
463impl<T: NodeStoreScan> NodeStoreScan for Arc<T> {
464 type Error = T::Error;
465
466 fn list_node_cids(&self) -> Result<Vec<Cid>, Self::Error> {
467 (**self).list_node_cids()
468 }
469}
470
471pub(crate) fn cid_from_store_key(key: &[u8], context: impl AsRef<str>) -> Result<Cid, String> {
472 let context = context.as_ref();
473 if key.len() != 32 {
474 return Err(format!(
475 "{context} key has invalid CID length {}, expected 32",
476 key.len()
477 ));
478 }
479 let mut bytes = [0u8; 32];
480 bytes.copy_from_slice(key);
481 Ok(Cid(bytes))
482}
483
484pub(crate) fn sort_cids(cids: &mut [Cid]) {
485 cids.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
486}
487
488#[allow(async_fn_in_trait)]
498pub trait AsyncStore {
499 type Error: std::error::Error + 'static;
501
502 async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error>;
504
505 async fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
507 self.get(key)
508 .await
509 .map(|value| value.map(|bytes| Arc::from(bytes.into_boxed_slice())))
510 }
511
512 async fn batch_get_shared_ordered_unique(
514 &self,
515 keys: &[&[u8]],
516 ) -> Result<SharedReadBatch, Self::Error> {
517 self.batch_get_ordered_unique(keys).await.map(|values| {
518 values
519 .into_iter()
520 .map(|value| value.map(|bytes| Arc::from(bytes.into_boxed_slice())))
521 .collect()
522 })
523 }
524
525 fn has_native_shared_reads(&self) -> bool {
527 false
528 }
529
530 async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error>;
532
533 async fn delete(&self, key: &[u8]) -> Result<(), Self::Error>;
535
536 async fn batch(&self, ops: &[BatchOp<'_>]) -> Result<(), Self::Error>;
538
539 async fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
544 let ordered = self.batch_get_ordered(keys).await?;
545 let mut results = HashMap::with_capacity(ordered.len());
546 for (key, value) in keys.iter().zip(ordered) {
547 if let Some(value) = value {
548 results.insert((*key).to_vec(), value);
549 }
550 }
551 Ok(results)
552 }
553
554 async fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
561 async_batch_get_ordered_with_limit(self, keys, self.read_parallelism()).await
562 }
563
564 async fn batch_get_ordered_unique(
570 &self,
571 keys: &[&[u8]],
572 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
573 async_batch_get_ordered_unique_with_limit(self, keys, self.read_parallelism()).await
574 }
575
576 fn prefers_batch_reads(&self) -> bool {
578 false
579 }
580
581 fn read_parallelism(&self) -> usize {
587 1
588 }
589
590 async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
592 let ops: Vec<BatchOp<'_>> = entries
593 .iter()
594 .map(|(key, value)| BatchOp::Upsert { key, value })
595 .collect();
596 self.batch(&ops).await
597 }
598
599 fn supports_hints(&self) -> bool {
601 false
602 }
603
604 fn prefers_rightmost_path_hints(&self) -> bool {
610 false
611 }
612
613 async fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
615 let _ = (namespace, key);
616 Ok(None)
617 }
618
619 async fn put_hint(
621 &self,
622 namespace: &[u8],
623 key: &[u8],
624 value: &[u8],
625 ) -> Result<(), Self::Error> {
626 let _ = (namespace, key, value);
627 Ok(())
628 }
629
630 async fn batch_put_with_hint(
632 &self,
633 entries: &[(&[u8], &[u8])],
634 namespace: &[u8],
635 key: &[u8],
636 value: &[u8],
637 ) -> Result<(), Self::Error> {
638 self.batch_put(entries).await?;
639 self.put_hint(namespace, key, value).await
640 }
641
642 #[inline(always)]
648 fn publish_nodes<'a>(
649 &'a self,
650 publication: NodePublication<'a>,
651 ) -> impl Future<Output = Result<(), Self::Error>> + 'a {
652 use futures_util::future::Either;
653
654 let entries = publication.entries;
655 match publication.hint {
656 Some(hint) => Either::Left(self.batch_put_with_hint(
657 entries,
658 hint.namespace(),
659 hint.key(),
660 hint.value(),
661 )),
662 None => Either::Right(self.batch_put(entries)),
663 }
664 }
665}
666async fn async_batch_get_ordered_with_limit<S: AsyncStore + ?Sized>(
667 store: &S,
668 keys: &[&[u8]],
669 max_in_flight: usize,
670) -> Result<Vec<Option<Vec<u8>>>, S::Error> {
671 if keys.len() < 2 {
672 return async_batch_get_ordered_unique_with_limit(store, keys, max_in_flight).await;
673 }
674
675 let plan = OrderedBatchReadPlan::new(keys);
676 let unique_values =
677 async_batch_get_ordered_unique_with_limit(store, plan.unique_keys(), max_in_flight).await?;
678 Ok(plan.expand_owned(unique_values))
679}
680async fn async_batch_get_ordered_unique_with_limit<S: AsyncStore + ?Sized>(
681 store: &S,
682 keys: &[&[u8]],
683 max_in_flight: usize,
684) -> Result<Vec<Option<Vec<u8>>>, S::Error> {
685 if keys.is_empty() {
686 return Ok(Vec::new());
687 }
688
689 let max_in_flight = max_in_flight.max(1);
690 if keys.len() < 2 || max_in_flight == 1 {
691 let mut values = Vec::with_capacity(keys.len());
692 for key in keys {
693 values.push(store.get(key).await?);
694 }
695 return Ok(values);
696 }
697
698 use futures_util::stream::{FuturesUnordered, StreamExt as _};
699
700 let mut values = vec![None; keys.len()];
701 let mut next_idx = 0;
702 let mut in_flight = FuturesUnordered::new();
703
704 while next_idx < keys.len() && in_flight.len() < max_in_flight {
705 in_flight.push(async_get_indexed(store, next_idx, keys[next_idx]));
706 next_idx += 1;
707 }
708
709 while let Some((idx, result)) = in_flight.next().await {
710 values[idx] = result?;
711
712 if next_idx < keys.len() {
713 in_flight.push(async_get_indexed(store, next_idx, keys[next_idx]));
714 next_idx += 1;
715 }
716 }
717
718 Ok(values)
719}
720async fn async_get_indexed<S: AsyncStore + ?Sized>(
721 store: &S,
722 idx: usize,
723 key: &[u8],
724) -> (usize, Result<Option<Vec<u8>>, S::Error>) {
725 (idx, store.get(key).await)
726}
727
728#[derive(Clone, Debug)]
743pub struct SyncStoreAsAsync<S> {
744 inner: S,
745}
746impl<S> SyncStoreAsAsync<S> {
747 pub fn new(inner: S) -> Self {
749 Self { inner }
750 }
751
752 pub fn inner(&self) -> &S {
754 &self.inner
755 }
756
757 pub fn into_inner(self) -> S {
759 self.inner
760 }
761
762 pub(crate) fn ready<F: Future>(&self, future: F) -> ReadyOnly<F> {
764 ready_only(future)
765 }
766}
767impl<S: Store> AsyncStore for SyncStoreAsAsync<S> {
768 type Error = S::Error;
769
770 async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
771 self.inner.get(key)
772 }
773
774 async fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
775 self.inner.get_shared(key)
776 }
777
778 async fn batch_get_shared_ordered_unique(
779 &self,
780 keys: &[&[u8]],
781 ) -> Result<SharedReadBatch, Self::Error> {
782 self.inner.batch_get_shared_ordered_unique(keys)
783 }
784
785 fn has_native_shared_reads(&self) -> bool {
786 self.inner.has_native_shared_reads()
787 }
788
789 async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
790 self.inner.put(key, value)
791 }
792
793 async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
794 self.inner.delete(key)
795 }
796
797 async fn batch(&self, ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
798 self.inner.batch(ops)
799 }
800
801 async fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
802 self.inner.batch_get(keys)
803 }
804
805 async fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
806 self.inner.batch_get_ordered(keys)
807 }
808
809 async fn batch_get_ordered_unique(
810 &self,
811 keys: &[&[u8]],
812 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
813 self.inner.batch_get_ordered_unique(keys)
814 }
815
816 fn prefers_batch_reads(&self) -> bool {
817 self.inner.prefers_batch_reads()
818 }
819
820 async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
821 self.inner.batch_put(entries)
822 }
823
824 fn supports_hints(&self) -> bool {
825 self.inner.supports_hints()
826 }
827
828 fn prefers_rightmost_path_hints(&self) -> bool {
829 self.inner.prefers_rightmost_path_hints()
830 }
831
832 async fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
833 self.inner.get_hint(namespace, key)
834 }
835
836 async fn put_hint(
837 &self,
838 namespace: &[u8],
839 key: &[u8],
840 value: &[u8],
841 ) -> Result<(), Self::Error> {
842 self.inner.put_hint(namespace, key, value)
843 }
844
845 async fn batch_put_with_hint(
846 &self,
847 entries: &[(&[u8], &[u8])],
848 namespace: &[u8],
849 key: &[u8],
850 value: &[u8],
851 ) -> Result<(), Self::Error> {
852 self.inner
853 .batch_put_with_hint(entries, namespace, key, value)
854 }
855
856 #[inline(always)]
857 async fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
858 self.inner.publish_nodes(publication)
859 }
860}
861impl<S: ManifestStore> AsyncManifestStore for SyncStoreAsAsync<S> {
862 type Error = S::Error;
863
864 async fn get_root(&self, name: &[u8]) -> Result<Option<RootManifest>, Self::Error> {
865 self.inner.get_root(name)
866 }
867
868 async fn put_root(&self, name: &[u8], manifest: &RootManifest) -> Result<(), Self::Error> {
869 self.inner.put_root(name, manifest)
870 }
871
872 async fn delete_root(&self, name: &[u8]) -> Result<(), Self::Error> {
873 self.inner.delete_root(name)
874 }
875
876 async fn compare_and_swap_root(
877 &self,
878 name: &[u8],
879 expected: Option<&RootManifest>,
880 new: Option<&RootManifest>,
881 ) -> Result<ManifestUpdate, Self::Error> {
882 self.inner.compare_and_swap_root(name, expected, new)
883 }
884}
885impl<S: ManifestStoreScan> AsyncManifestStoreScan for SyncStoreAsAsync<S> {
886 async fn list_roots(&self) -> Result<Vec<NamedRootManifest>, Self::Error> {
887 self.inner.list_roots()
888 }
889}
890
891#[cfg(feature = "tokio")]
893#[derive(Debug)]
894pub enum TokioBlockingStoreError<E> {
895 Store(E),
897 Join(tokio::task::JoinError),
900}
901
902#[cfg(feature = "tokio")]
903impl<E: std::fmt::Display> std::fmt::Display for TokioBlockingStoreError<E> {
904 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
905 match self {
906 Self::Store(err) => write!(f, "store error: {err}"),
907 Self::Join(err) => write!(f, "tokio blocking task failed: {err}"),
908 }
909 }
910}
911
912#[cfg(feature = "tokio")]
913impl<E> std::error::Error for TokioBlockingStoreError<E>
914where
915 E: std::error::Error + 'static,
916{
917 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
918 match self {
919 Self::Store(err) => Some(err),
920 Self::Join(err) => Some(err),
921 }
922 }
923}
924
925#[cfg(feature = "tokio")]
932#[derive(Debug)]
933pub struct TokioBlockingStore<S> {
934 inner: std::sync::Arc<S>,
935}
936
937#[cfg(feature = "tokio")]
938impl<S> Clone for TokioBlockingStore<S> {
939 fn clone(&self) -> Self {
940 Self {
941 inner: self.inner.clone(),
942 }
943 }
944}
945
946#[cfg(feature = "tokio")]
947impl<S> TokioBlockingStore<S> {
948 pub fn new(inner: S) -> Self {
950 Self {
951 inner: std::sync::Arc::new(inner),
952 }
953 }
954
955 pub fn from_arc(inner: std::sync::Arc<S>) -> Self {
957 Self { inner }
958 }
959
960 pub fn inner(&self) -> &S {
962 &self.inner
963 }
964
965 pub fn shared(&self) -> std::sync::Arc<S> {
967 self.inner.clone()
968 }
969}
970
971#[cfg(feature = "tokio")]
972async fn spawn_store_blocking<S, F, R>(
973 store: std::sync::Arc<S>,
974 operation: F,
975) -> Result<R, TokioBlockingStoreError<S::Error>>
976where
977 S: Store + 'static,
978 F: FnOnce(std::sync::Arc<S>) -> Result<R, S::Error> + Send + 'static,
979 R: Send + 'static,
980{
981 tokio::task::spawn_blocking(move || operation(store))
982 .await
983 .map_err(TokioBlockingStoreError::Join)?
984 .map_err(TokioBlockingStoreError::Store)
985}
986
987#[cfg(feature = "tokio")]
988impl<S> AsyncStore for TokioBlockingStore<S>
989where
990 S: Store + 'static,
991{
992 type Error = TokioBlockingStoreError<S::Error>;
993
994 async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
995 let key = key.to_vec();
996 spawn_store_blocking(self.inner.clone(), move |store| store.get(&key)).await
997 }
998
999 async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1000 let key = key.to_vec();
1001 let value = value.to_vec();
1002 spawn_store_blocking(self.inner.clone(), move |store| store.put(&key, &value)).await
1003 }
1004
1005 async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1006 let key = key.to_vec();
1007 spawn_store_blocking(self.inner.clone(), move |store| store.delete(&key)).await
1008 }
1009
1010 async fn batch(&self, ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1011 let owned_ops = ops
1012 .iter()
1013 .map(|op| match op {
1014 BatchOp::Upsert { key, value } => (true, key.to_vec(), value.to_vec()),
1015 BatchOp::Delete { key } => (false, key.to_vec(), Vec::new()),
1016 })
1017 .collect::<Vec<_>>();
1018
1019 spawn_store_blocking(self.inner.clone(), move |store| {
1020 let ops = owned_ops
1021 .iter()
1022 .map(|(is_upsert, key, value)| {
1023 if *is_upsert {
1024 BatchOp::Upsert {
1025 key: key.as_slice(),
1026 value: value.as_slice(),
1027 }
1028 } else {
1029 BatchOp::Delete {
1030 key: key.as_slice(),
1031 }
1032 }
1033 })
1034 .collect::<Vec<_>>();
1035 store.batch(&ops)
1036 })
1037 .await
1038 }
1039
1040 async fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
1041 let keys = keys.iter().map(|key| key.to_vec()).collect::<Vec<_>>();
1042 spawn_store_blocking(self.inner.clone(), move |store| {
1043 let keys = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
1044 store.batch_get(&keys)
1045 })
1046 .await
1047 }
1048
1049 async fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1050 let keys = keys.iter().map(|key| key.to_vec()).collect::<Vec<_>>();
1051 spawn_store_blocking(self.inner.clone(), move |store| {
1052 let keys = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
1053 store.batch_get_ordered(&keys)
1054 })
1055 .await
1056 }
1057
1058 async fn batch_get_ordered_unique(
1059 &self,
1060 keys: &[&[u8]],
1061 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1062 let keys = keys.iter().map(|key| key.to_vec()).collect::<Vec<_>>();
1063 spawn_store_blocking(self.inner.clone(), move |store| {
1064 let keys = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
1065 store.batch_get_ordered_unique(&keys)
1066 })
1067 .await
1068 }
1069
1070 fn prefers_batch_reads(&self) -> bool {
1071 self.inner.prefers_batch_reads()
1072 }
1073
1074 async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1075 let entries = entries
1076 .iter()
1077 .map(|(key, value)| (key.to_vec(), value.to_vec()))
1078 .collect::<Vec<_>>();
1079 spawn_store_blocking(self.inner.clone(), move |store| {
1080 let entries = entries
1081 .iter()
1082 .map(|(key, value)| (key.as_slice(), value.as_slice()))
1083 .collect::<Vec<_>>();
1084 store.batch_put(&entries)
1085 })
1086 .await
1087 }
1088
1089 fn supports_hints(&self) -> bool {
1090 self.inner.supports_hints()
1091 }
1092
1093 fn prefers_rightmost_path_hints(&self) -> bool {
1094 self.inner.prefers_rightmost_path_hints()
1095 }
1096
1097 async fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1098 let namespace = namespace.to_vec();
1099 let key = key.to_vec();
1100 spawn_store_blocking(self.inner.clone(), move |store| {
1101 store.get_hint(&namespace, &key)
1102 })
1103 .await
1104 }
1105
1106 async fn put_hint(
1107 &self,
1108 namespace: &[u8],
1109 key: &[u8],
1110 value: &[u8],
1111 ) -> Result<(), Self::Error> {
1112 let namespace = namespace.to_vec();
1113 let key = key.to_vec();
1114 let value = value.to_vec();
1115 spawn_store_blocking(self.inner.clone(), move |store| {
1116 store.put_hint(&namespace, &key, &value)
1117 })
1118 .await
1119 }
1120
1121 async fn batch_put_with_hint(
1122 &self,
1123 entries: &[(&[u8], &[u8])],
1124 namespace: &[u8],
1125 key: &[u8],
1126 value: &[u8],
1127 ) -> Result<(), Self::Error> {
1128 let entries = entries
1129 .iter()
1130 .map(|(key, value)| (key.to_vec(), value.to_vec()))
1131 .collect::<Vec<_>>();
1132 let namespace = namespace.to_vec();
1133 let key = key.to_vec();
1134 let value = value.to_vec();
1135 spawn_store_blocking(self.inner.clone(), move |store| {
1136 let entries = entries
1137 .iter()
1138 .map(|(key, value)| (key.as_slice(), value.as_slice()))
1139 .collect::<Vec<_>>();
1140 store.batch_put_with_hint(&entries, &namespace, &key, &value)
1141 })
1142 .await
1143 }
1144
1145 async fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1146 let entries = publication
1147 .entries()
1148 .iter()
1149 .map(|(key, value)| (key.to_vec(), value.to_vec()))
1150 .collect::<Vec<_>>();
1151 let hint = publication.hint().map(|hint| {
1152 (
1153 hint.namespace().to_vec(),
1154 hint.key().to_vec(),
1155 hint.value().to_vec(),
1156 )
1157 });
1158 let origin = publication.origin();
1159
1160 spawn_store_blocking(self.inner.clone(), move |store| {
1161 let entries = entries
1162 .iter()
1163 .map(|(key, value)| (key.as_slice(), value.as_slice()))
1164 .collect::<Vec<_>>();
1165 let publication = match hint.as_ref() {
1166 Some((namespace, key, value)) => NodePublication::with_hint(
1167 &entries,
1168 NodePublicationHint::new(namespace, key, value),
1169 origin,
1170 ),
1171 None => NodePublication::new(&entries, origin),
1172 };
1173 store.publish_nodes(publication)
1174 })
1175 .await
1176 }
1177}
1178
1179#[cfg(feature = "tokio")]
1180async fn spawn_manifest_blocking<S, F, R>(
1181 store: std::sync::Arc<S>,
1182 operation: F,
1183) -> Result<R, TokioBlockingStoreError<<S as ManifestStore>::Error>>
1184where
1185 S: ManifestStore + 'static,
1186 F: FnOnce(std::sync::Arc<S>) -> Result<R, <S as ManifestStore>::Error> + Send + 'static,
1187 R: Send + 'static,
1188{
1189 tokio::task::spawn_blocking(move || operation(store))
1190 .await
1191 .map_err(TokioBlockingStoreError::Join)?
1192 .map_err(TokioBlockingStoreError::Store)
1193}
1194
1195#[cfg(feature = "tokio")]
1196impl<S> AsyncManifestStore for TokioBlockingStore<S>
1197where
1198 S: ManifestStore + 'static,
1199{
1200 type Error = TokioBlockingStoreError<<S as ManifestStore>::Error>;
1201
1202 async fn get_root(&self, name: &[u8]) -> Result<Option<RootManifest>, Self::Error> {
1203 let name = name.to_vec();
1204 spawn_manifest_blocking(self.inner.clone(), move |store| store.get_root(&name)).await
1205 }
1206
1207 async fn put_root(&self, name: &[u8], manifest: &RootManifest) -> Result<(), Self::Error> {
1208 let name = name.to_vec();
1209 let manifest = manifest.clone();
1210 spawn_manifest_blocking(self.inner.clone(), move |store| {
1211 store.put_root(&name, &manifest)
1212 })
1213 .await
1214 }
1215
1216 async fn delete_root(&self, name: &[u8]) -> Result<(), Self::Error> {
1217 let name = name.to_vec();
1218 spawn_manifest_blocking(self.inner.clone(), move |store| store.delete_root(&name)).await
1219 }
1220
1221 async fn compare_and_swap_root(
1222 &self,
1223 name: &[u8],
1224 expected: Option<&RootManifest>,
1225 new: Option<&RootManifest>,
1226 ) -> Result<ManifestUpdate, Self::Error> {
1227 let name = name.to_vec();
1228 let expected = expected.cloned();
1229 let new = new.cloned();
1230 spawn_manifest_blocking(self.inner.clone(), move |store| {
1231 store.compare_and_swap_root(&name, expected.as_ref(), new.as_ref())
1232 })
1233 .await
1234 }
1235}
1236
1237#[cfg(feature = "tokio")]
1238impl<S> AsyncManifestStoreScan for TokioBlockingStore<S>
1239where
1240 S: ManifestStoreScan + 'static,
1241{
1242 async fn list_roots(&self) -> Result<Vec<NamedRootManifest>, Self::Error> {
1243 spawn_manifest_blocking(self.inner.clone(), move |store| store.list_roots()).await
1244 }
1245}
1246
1247impl<T: Store> Store for std::sync::Arc<T> {
1250 type Error = T::Error;
1251
1252 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1253 (**self).get(key)
1254 }
1255
1256 fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
1257 (**self).get_shared(key)
1258 }
1259
1260 fn batch_get_shared_ordered_unique(
1261 &self,
1262 keys: &[&[u8]],
1263 ) -> Result<SharedReadBatch, Self::Error> {
1264 (**self).batch_get_shared_ordered_unique(keys)
1265 }
1266
1267 fn has_native_shared_reads(&self) -> bool {
1268 (**self).has_native_shared_reads()
1269 }
1270
1271 fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1272 (**self).put(key, value)
1273 }
1274
1275 fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1276 (**self).delete(key)
1277 }
1278
1279 fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
1280 (**self).batch(ops)
1281 }
1282
1283 fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
1284 (**self).batch_get(keys)
1285 }
1286
1287 fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1288 (**self).batch_get_ordered(keys)
1289 }
1290
1291 fn batch_get_ordered_unique(
1292 &self,
1293 keys: &[&[u8]],
1294 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1295 (**self).batch_get_ordered_unique(keys)
1296 }
1297
1298 fn prefers_batch_reads(&self) -> bool {
1299 (**self).prefers_batch_reads()
1300 }
1301
1302 fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1303 (**self).batch_put(entries)
1304 }
1305
1306 fn supports_hints(&self) -> bool {
1307 (**self).supports_hints()
1308 }
1309
1310 fn prefers_rightmost_path_hints(&self) -> bool {
1311 (**self).prefers_rightmost_path_hints()
1312 }
1313
1314 fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1315 (**self).get_hint(namespace, key)
1316 }
1317
1318 fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1319 (**self).put_hint(namespace, key, value)
1320 }
1321
1322 fn batch_put_with_hint(
1323 &self,
1324 entries: &[(&[u8], &[u8])],
1325 namespace: &[u8],
1326 key: &[u8],
1327 value: &[u8],
1328 ) -> Result<(), Self::Error> {
1329 (**self).batch_put_with_hint(entries, namespace, key, value)
1330 }
1331
1332 fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1333 (**self).publish_nodes(publication)
1334 }
1335}
1336
1337impl<T: Store + ?Sized> Store for &T {
1343 type Error = T::Error;
1344
1345 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1346 (**self).get(key)
1347 }
1348
1349 fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
1350 (**self).get_shared(key)
1351 }
1352
1353 fn batch_get_shared_ordered_unique(
1354 &self,
1355 keys: &[&[u8]],
1356 ) -> Result<SharedReadBatch, Self::Error> {
1357 (**self).batch_get_shared_ordered_unique(keys)
1358 }
1359
1360 fn has_native_shared_reads(&self) -> bool {
1361 (**self).has_native_shared_reads()
1362 }
1363
1364 fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1365 (**self).put(key, value)
1366 }
1367
1368 fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1369 (**self).delete(key)
1370 }
1371
1372 fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
1373 (**self).batch(ops)
1374 }
1375
1376 fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
1377 (**self).batch_get(keys)
1378 }
1379
1380 fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1381 (**self).batch_get_ordered(keys)
1382 }
1383
1384 fn batch_get_ordered_unique(
1385 &self,
1386 keys: &[&[u8]],
1387 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1388 (**self).batch_get_ordered_unique(keys)
1389 }
1390
1391 fn prefers_batch_reads(&self) -> bool {
1392 (**self).prefers_batch_reads()
1393 }
1394
1395 fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1396 (**self).batch_put(entries)
1397 }
1398
1399 fn supports_hints(&self) -> bool {
1400 (**self).supports_hints()
1401 }
1402
1403 fn prefers_rightmost_path_hints(&self) -> bool {
1404 (**self).prefers_rightmost_path_hints()
1405 }
1406
1407 fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1408 (**self).get_hint(namespace, key)
1409 }
1410
1411 fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1412 (**self).put_hint(namespace, key, value)
1413 }
1414
1415 fn batch_put_with_hint(
1416 &self,
1417 entries: &[(&[u8], &[u8])],
1418 namespace: &[u8],
1419 key: &[u8],
1420 value: &[u8],
1421 ) -> Result<(), Self::Error> {
1422 (**self).batch_put_with_hint(entries, namespace, key, value)
1423 }
1424
1425 fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1426 (**self).publish_nodes(publication)
1427 }
1428}
1429impl<T: AsyncStore> AsyncStore for std::sync::Arc<T> {
1430 type Error = T::Error;
1431
1432 async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1433 (**self).get(key).await
1434 }
1435
1436 async fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
1437 (**self).get_shared(key).await
1438 }
1439
1440 async fn batch_get_shared_ordered_unique(
1441 &self,
1442 keys: &[&[u8]],
1443 ) -> Result<SharedReadBatch, Self::Error> {
1444 (**self).batch_get_shared_ordered_unique(keys).await
1445 }
1446
1447 fn has_native_shared_reads(&self) -> bool {
1448 (**self).has_native_shared_reads()
1449 }
1450
1451 async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1452 (**self).put(key, value).await
1453 }
1454
1455 async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1456 (**self).delete(key).await
1457 }
1458
1459 async fn batch(&self, ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1460 (**self).batch(ops).await
1461 }
1462
1463 async fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
1464 (**self).batch_get(keys).await
1465 }
1466
1467 async fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1468 (**self).batch_get_ordered(keys).await
1469 }
1470
1471 async fn batch_get_ordered_unique(
1472 &self,
1473 keys: &[&[u8]],
1474 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1475 (**self).batch_get_ordered_unique(keys).await
1476 }
1477
1478 fn prefers_batch_reads(&self) -> bool {
1479 (**self).prefers_batch_reads()
1480 }
1481
1482 fn read_parallelism(&self) -> usize {
1483 (**self).read_parallelism()
1484 }
1485
1486 async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1487 (**self).batch_put(entries).await
1488 }
1489
1490 fn supports_hints(&self) -> bool {
1491 (**self).supports_hints()
1492 }
1493
1494 fn prefers_rightmost_path_hints(&self) -> bool {
1495 (**self).prefers_rightmost_path_hints()
1496 }
1497
1498 async fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1499 (**self).get_hint(namespace, key).await
1500 }
1501
1502 async fn put_hint(
1503 &self,
1504 namespace: &[u8],
1505 key: &[u8],
1506 value: &[u8],
1507 ) -> Result<(), Self::Error> {
1508 (**self).put_hint(namespace, key, value).await
1509 }
1510
1511 async fn batch_put_with_hint(
1512 &self,
1513 entries: &[(&[u8], &[u8])],
1514 namespace: &[u8],
1515 key: &[u8],
1516 value: &[u8],
1517 ) -> Result<(), Self::Error> {
1518 (**self)
1519 .batch_put_with_hint(entries, namespace, key, value)
1520 .await
1521 }
1522
1523 async fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1524 (**self).publish_nodes(publication).await
1525 }
1526}
1527
1528#[cfg(test)]
1529mod tests {
1530 use super::*;
1531 use std::collections::BTreeMap;
1532 use std::sync::atomic::{AtomicUsize, Ordering};
1533 use std::sync::Mutex;
1534 use std::{
1535 future::Future,
1536 pin::Pin,
1537 task::{Context, Poll},
1538 };
1539
1540 #[derive(Debug)]
1541 struct DefaultReadStoreError;
1542
1543 impl std::fmt::Display for DefaultReadStoreError {
1544 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1545 f.write_str("default read store error")
1546 }
1547 }
1548
1549 impl std::error::Error for DefaultReadStoreError {}
1550
1551 #[derive(Default)]
1552 struct DefaultReadStore {
1553 data: Mutex<BTreeMap<Vec<u8>, Vec<u8>>>,
1554 get_calls: AtomicUsize,
1555 }
1556
1557 impl DefaultReadStore {
1558 fn with_entries(entries: &[(&[u8], &[u8])]) -> Self {
1559 let mut data = BTreeMap::new();
1560 for (key, value) in entries {
1561 data.insert(key.to_vec(), value.to_vec());
1562 }
1563
1564 Self {
1565 data: Mutex::new(data),
1566 get_calls: AtomicUsize::new(0),
1567 }
1568 }
1569 }
1570
1571 impl Store for DefaultReadStore {
1572 type Error = DefaultReadStoreError;
1573
1574 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1575 self.get_calls.fetch_add(1, Ordering::Relaxed);
1576 Ok(self.data.lock().unwrap().get(key).cloned())
1577 }
1578
1579 fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1580 self.data
1581 .lock()
1582 .unwrap()
1583 .insert(key.to_vec(), value.to_vec());
1584 Ok(())
1585 }
1586
1587 fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1588 self.data.lock().unwrap().remove(key);
1589 Ok(())
1590 }
1591
1592 fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
1593 let mut data = self.data.lock().unwrap();
1594 for op in ops {
1595 match op {
1596 BatchOp::Upsert { key, value } => {
1597 data.insert(key.to_vec(), value.to_vec());
1598 }
1599 BatchOp::Delete { key } => {
1600 data.remove(*key);
1601 }
1602 }
1603 }
1604 Ok(())
1605 }
1606 }
1607
1608 type OwnedNodeEntry = (Vec<u8>, Vec<u8>);
1609 type OwnedPublication = Vec<OwnedNodeEntry>;
1610 type OwnedPublicationHint = (Vec<u8>, Vec<u8>, Vec<u8>);
1611
1612 #[derive(Default)]
1613 struct DefaultPublicationStore {
1614 batch_calls: AtomicUsize,
1615 hinted_batch_calls: AtomicUsize,
1616 published_entries: Mutex<Vec<OwnedPublication>>,
1617 last_hint: Mutex<Option<OwnedPublicationHint>>,
1618 }
1619
1620 impl DefaultPublicationStore {
1621 fn record_entries(&self, entries: &[(&[u8], &[u8])]) {
1622 self.published_entries.lock().unwrap().push(
1623 entries
1624 .iter()
1625 .map(|(key, value)| (key.to_vec(), value.to_vec()))
1626 .collect(),
1627 );
1628 }
1629 }
1630
1631 impl Store for DefaultPublicationStore {
1632 type Error = DefaultReadStoreError;
1633
1634 fn get(&self, _key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1635 Ok(None)
1636 }
1637
1638 fn put(&self, _key: &[u8], _value: &[u8]) -> Result<(), Self::Error> {
1639 Ok(())
1640 }
1641
1642 fn delete(&self, _key: &[u8]) -> Result<(), Self::Error> {
1643 Ok(())
1644 }
1645
1646 fn batch(&self, _ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1647 Ok(())
1648 }
1649
1650 fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1651 self.batch_calls.fetch_add(1, Ordering::Relaxed);
1652 self.record_entries(entries);
1653 Ok(())
1654 }
1655
1656 fn batch_put_with_hint(
1657 &self,
1658 entries: &[(&[u8], &[u8])],
1659 namespace: &[u8],
1660 key: &[u8],
1661 value: &[u8],
1662 ) -> Result<(), Self::Error> {
1663 self.hinted_batch_calls.fetch_add(1, Ordering::Relaxed);
1664 self.record_entries(entries);
1665 *self.last_hint.lock().unwrap() =
1666 Some((namespace.to_vec(), key.to_vec(), value.to_vec()));
1667 Ok(())
1668 }
1669 }
1670
1671 #[derive(Default)]
1672 struct ForwardingPublicationStore {
1673 publication_calls: AtomicUsize,
1674 last_origin: Mutex<Option<PublicationOrigin>>,
1675 }
1676
1677 impl ForwardingPublicationStore {
1678 fn record_publication(&self, publication: NodePublication<'_>) {
1679 self.publication_calls.fetch_add(1, Ordering::Relaxed);
1680 *self.last_origin.lock().unwrap() = Some(publication.origin());
1681 }
1682 }
1683
1684 impl Store for ForwardingPublicationStore {
1685 type Error = DefaultReadStoreError;
1686
1687 fn get(&self, _key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1688 Ok(None)
1689 }
1690
1691 fn put(&self, _key: &[u8], _value: &[u8]) -> Result<(), Self::Error> {
1692 Ok(())
1693 }
1694
1695 fn delete(&self, _key: &[u8]) -> Result<(), Self::Error> {
1696 Ok(())
1697 }
1698
1699 fn batch(&self, _ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1700 Ok(())
1701 }
1702
1703 fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1704 self.record_publication(publication);
1705 Ok(())
1706 }
1707 }
1708
1709 impl AsyncStore for ForwardingPublicationStore {
1710 type Error = DefaultReadStoreError;
1711
1712 async fn get(&self, _key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1713 Ok(None)
1714 }
1715
1716 async fn put(&self, _key: &[u8], _value: &[u8]) -> Result<(), Self::Error> {
1717 Ok(())
1718 }
1719
1720 async fn delete(&self, _key: &[u8]) -> Result<(), Self::Error> {
1721 Ok(())
1722 }
1723
1724 async fn batch(&self, _ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1725 Ok(())
1726 }
1727
1728 async fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
1729 self.record_publication(publication);
1730 Ok(())
1731 }
1732 }
1733
1734 fn block_on<F: Future>(future: F) -> F::Output {
1735 let waker = futures_util::task::noop_waker();
1736 let mut cx = Context::from_waker(&waker);
1737 let mut future = Box::pin(future);
1738
1739 loop {
1740 match future.as_mut().poll(&mut cx) {
1741 Poll::Ready(value) => return value,
1742 Poll::Pending => std::thread::yield_now(),
1743 }
1744 }
1745 }
1746 struct YieldOnce {
1747 yielded: bool,
1748 }
1749 impl Future for YieldOnce {
1750 type Output = ();
1751
1752 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1753 if self.yielded {
1754 Poll::Ready(())
1755 } else {
1756 self.yielded = true;
1757 cx.waker().wake_by_ref();
1758 Poll::Pending
1759 }
1760 }
1761 }
1762 struct DefaultAsyncReadStore {
1763 data: Mutex<BTreeMap<Vec<u8>, Vec<u8>>>,
1764 get_calls: AtomicUsize,
1765 in_flight: AtomicUsize,
1766 max_in_flight: AtomicUsize,
1767 read_parallelism: usize,
1768 }
1769 impl DefaultAsyncReadStore {
1770 fn with_entries(read_parallelism: usize, entries: &[(&[u8], &[u8])]) -> Self {
1771 let mut data = BTreeMap::new();
1772 for (key, value) in entries {
1773 data.insert(key.to_vec(), value.to_vec());
1774 }
1775
1776 Self {
1777 data: Mutex::new(data),
1778 get_calls: AtomicUsize::new(0),
1779 in_flight: AtomicUsize::new(0),
1780 max_in_flight: AtomicUsize::new(0),
1781 read_parallelism,
1782 }
1783 }
1784 }
1785 impl AsyncStore for DefaultAsyncReadStore {
1786 type Error = DefaultReadStoreError;
1787
1788 async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1789 self.get_calls.fetch_add(1, Ordering::Relaxed);
1790 let current = self.in_flight.fetch_add(1, Ordering::Relaxed) + 1;
1791 self.max_in_flight.fetch_max(current, Ordering::Relaxed);
1792
1793 YieldOnce { yielded: false }.await;
1794
1795 let value = self.data.lock().unwrap().get(key).cloned();
1796 self.in_flight.fetch_sub(1, Ordering::Relaxed);
1797 Ok(value)
1798 }
1799
1800 async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1801 self.data
1802 .lock()
1803 .unwrap()
1804 .insert(key.to_vec(), value.to_vec());
1805 Ok(())
1806 }
1807
1808 async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1809 self.data.lock().unwrap().remove(key);
1810 Ok(())
1811 }
1812
1813 async fn batch(&self, ops: &[BatchOp<'_>]) -> Result<(), Self::Error> {
1814 let mut data = self.data.lock().unwrap();
1815 for op in ops {
1816 match op {
1817 BatchOp::Upsert { key, value } => {
1818 data.insert(key.to_vec(), value.to_vec());
1819 }
1820 BatchOp::Delete { key } => {
1821 data.remove(*key);
1822 }
1823 }
1824 }
1825 Ok(())
1826 }
1827
1828 fn read_parallelism(&self) -> usize {
1829 self.read_parallelism
1830 }
1831 }
1832
1833 #[test]
1834 fn ordered_batch_read_plan_keeps_unique_batches_identity() {
1835 let keys: Vec<&[u8]> = vec![b"a", b"b", b"missing"];
1836 let plan = OrderedBatchReadPlan::new(&keys);
1837
1838 assert!(plan.is_identity());
1839 assert_eq!(
1840 plan.unique_keys(),
1841 &[b"a".as_slice(), b"b".as_slice(), b"missing".as_slice()]
1842 );
1843
1844 let values = vec![Some(b"1".to_vec()), Some(b"2".to_vec()), None];
1845 let values_ptr = values.as_ptr();
1846 let expanded = plan.expand_owned(values);
1847
1848 assert_eq!(expanded.as_ptr(), values_ptr);
1849 assert_eq!(
1850 expanded,
1851 vec![Some(b"1".to_vec()), Some(b"2".to_vec()), None]
1852 );
1853 }
1854
1855 #[test]
1856 fn publication_defaults_dispatch_once_and_preserve_borrowed_context() {
1857 let store = DefaultPublicationStore::default();
1858 let entries = [(b"node".as_slice(), b"bytes".as_slice())];
1859 let hint = NodePublicationHint::new(b"rightmost", b"root", b"path");
1860
1861 store
1862 .publish_nodes(NodePublication::new(
1863 &entries,
1864 PublicationOrigin::PointUpsert,
1865 ))
1866 .unwrap();
1867 store
1868 .publish_nodes(NodePublication::with_hint(
1869 &entries,
1870 hint,
1871 PublicationOrigin::TreeBuild,
1872 ))
1873 .unwrap();
1874
1875 assert_eq!(store.batch_calls.load(Ordering::Relaxed), 1);
1876 assert_eq!(store.hinted_batch_calls.load(Ordering::Relaxed), 1);
1877 assert_eq!(
1878 *store.last_hint.lock().unwrap(),
1879 Some((b"rightmost".to_vec(), b"root".to_vec(), b"path".to_vec()))
1880 );
1881 assert_eq!(
1882 *store.published_entries.lock().unwrap(),
1883 vec![
1884 vec![(b"node".to_vec(), b"bytes".to_vec())],
1885 vec![(b"node".to_vec(), b"bytes".to_vec())]
1886 ]
1887 );
1888 assert_eq!(hint.namespace(), b"rightmost");
1889 assert_eq!(hint.key(), b"root");
1890 assert_eq!(hint.value(), b"path");
1891 assert!(
1892 std::mem::size_of::<NodePublication<'static>>() <= std::mem::size_of::<[usize; 10]>()
1893 );
1894 }
1895
1896 #[test]
1897 fn sync_store_as_async_publication_is_ready_on_first_poll() {
1898 let store = SyncStoreAsAsync::new(DefaultPublicationStore::default());
1899 let entries = [(b"node".as_slice(), b"bytes".as_slice())];
1900 let publication = NodePublication::new(&entries, PublicationOrigin::PointDelete);
1901 let waker = futures_util::task::noop_waker();
1902 let mut context = Context::from_waker(&waker);
1903 let mut future = Box::pin(AsyncStore::publish_nodes(&store, publication));
1904
1905 assert!(matches!(
1906 future.as_mut().poll(&mut context),
1907 Poll::Ready(Ok(()))
1908 ));
1909 assert_eq!(store.inner().batch_calls.load(Ordering::Relaxed), 1);
1910 }
1911
1912 #[test]
1913 fn sync_reference_and_arc_publications_forward_exactly_once() {
1914 let store = Arc::new(ForwardingPublicationStore::default());
1915 let entries = [(b"node".as_slice(), b"bytes".as_slice())];
1916
1917 <Arc<ForwardingPublicationStore> as Store>::publish_nodes(
1918 &store,
1919 NodePublication::new(&entries, PublicationOrigin::Replication),
1920 )
1921 .unwrap();
1922 let borrowed = store.as_ref();
1923 <&ForwardingPublicationStore as Store>::publish_nodes(
1924 &borrowed,
1925 NodePublication::new(&entries, PublicationOrigin::Maintenance),
1926 )
1927 .unwrap();
1928
1929 assert_eq!(store.publication_calls.load(Ordering::Relaxed), 2);
1930 assert_eq!(
1931 *store.last_origin.lock().unwrap(),
1932 Some(PublicationOrigin::Maintenance)
1933 );
1934 }
1935
1936 #[test]
1937 fn async_arc_publication_forwards_exactly_once() {
1938 let store = Arc::new(ForwardingPublicationStore::default());
1939 let entries = [(b"node".as_slice(), b"bytes".as_slice())];
1940
1941 block_on(
1942 <Arc<ForwardingPublicationStore> as AsyncStore>::publish_nodes(
1943 &store,
1944 NodePublication::new(&entries, PublicationOrigin::BatchMutation),
1945 ),
1946 )
1947 .unwrap();
1948
1949 assert_eq!(store.publication_calls.load(Ordering::Relaxed), 1);
1950 assert_eq!(
1951 *store.last_origin.lock().unwrap(),
1952 Some(PublicationOrigin::BatchMutation)
1953 );
1954 }
1955
1956 #[test]
1957 fn ordered_batch_read_plan_deduplicates_and_expands_slots() {
1958 let keys: Vec<&[u8]> = vec![b"c", b"a", b"c", b"missing", b"missing", b"a"];
1959 let plan = OrderedBatchReadPlan::new(&keys);
1960
1961 assert!(!plan.is_identity());
1962 assert_eq!(
1963 plan.unique_keys(),
1964 &[b"c".as_slice(), b"a".as_slice(), b"missing".as_slice()]
1965 );
1966 assert_eq!(
1967 plan.expand(&[Some(b"3".to_vec()), Some(b"1".to_vec()), None]),
1968 vec![
1969 Some(b"3".to_vec()),
1970 Some(b"1".to_vec()),
1971 Some(b"3".to_vec()),
1972 None,
1973 None,
1974 Some(b"1".to_vec())
1975 ]
1976 );
1977 assert_eq!(
1978 plan.expand_owned(vec![Some(b"3".to_vec()), Some(b"1".to_vec()), None]),
1979 vec![
1980 Some(b"3".to_vec()),
1981 Some(b"1".to_vec()),
1982 Some(b"3".to_vec()),
1983 None,
1984 None,
1985 Some(b"1".to_vec())
1986 ]
1987 );
1988 }
1989
1990 #[test]
1991 fn default_batch_get_deduplicates_duplicate_keys() {
1992 let store = DefaultReadStore::with_entries(&[(b"a", b"1"), (b"b", b"2")]);
1993 let keys: Vec<&[u8]> = vec![b"a", b"a", b"missing", b"missing", b"b"];
1994
1995 let values = store.batch_get(&keys).unwrap();
1996
1997 assert_eq!(values.get(b"a".as_slice()), Some(&b"1".to_vec()));
1998 assert_eq!(values.get(b"b".as_slice()), Some(&b"2".to_vec()));
1999 assert!(!values.contains_key(b"missing".as_slice()));
2000 assert_eq!(
2001 store.get_calls.load(Ordering::Relaxed),
2002 3,
2003 "default batch_get should point-read each unique key at most once"
2004 );
2005 }
2006
2007 #[test]
2008 fn default_batch_get_ordered_deduplicates_while_preserving_slots() {
2009 let store = DefaultReadStore::with_entries(&[(b"a", b"1"), (b"b", b"2")]);
2010 let keys: Vec<&[u8]> = vec![b"a", b"a", b"missing", b"missing", b"b"];
2011
2012 let values = store.batch_get_ordered(&keys).unwrap();
2013
2014 assert_eq!(
2015 values,
2016 vec![
2017 Some(b"1".to_vec()),
2018 Some(b"1".to_vec()),
2019 None,
2020 None,
2021 Some(b"2".to_vec())
2022 ]
2023 );
2024 assert_eq!(
2025 store.get_calls.load(Ordering::Relaxed),
2026 3,
2027 "default ordered batch reads should preserve duplicate result slots without duplicate point reads"
2028 );
2029 }
2030
2031 #[test]
2032 fn default_unique_ordered_batch_reads_preserve_order_with_point_reads() {
2033 let store = DefaultReadStore::with_entries(&[(b"a", b"1"), (b"b", b"2")]);
2034 let keys: Vec<&[u8]> = vec![b"b", b"missing", b"a"];
2035
2036 let values = store.batch_get_ordered_unique(&keys).unwrap();
2037
2038 assert_eq!(values, vec![Some(b"2".to_vec()), None, Some(b"1".to_vec())]);
2039 assert_eq!(
2040 store.get_calls.load(Ordering::Relaxed),
2041 3,
2042 "unique ordered batch reads for point-read stores should read each requested key once"
2043 );
2044 }
2045 #[test]
2046 fn async_sync_store_adapter_preserves_default_store_behavior() {
2047 let store = DefaultReadStore::with_entries(&[(b"a", b"1"), (b"b", b"2")]);
2048 let store = SyncStoreAsAsync::new(store);
2049
2050 block_on(async {
2051 store.put(b"c", b"3").await.unwrap();
2052 store
2053 .batch_put(&[(b"d".as_slice(), b"4".as_slice())])
2054 .await
2055 .unwrap();
2056 store.delete(b"b").await.unwrap();
2057
2058 let keys: Vec<&[u8]> = vec![b"a", b"a", b"b", b"c", b"d"];
2059 let values = store.batch_get_ordered(&keys).await.unwrap();
2060 assert_eq!(
2061 values,
2062 vec![
2063 Some(b"1".to_vec()),
2064 Some(b"1".to_vec()),
2065 None,
2066 Some(b"3".to_vec()),
2067 Some(b"4".to_vec())
2068 ]
2069 );
2070
2071 let mapped = store.batch_get(&keys).await.unwrap();
2072 assert_eq!(mapped.get(b"a".as_slice()), Some(&b"1".to_vec()));
2073 assert_eq!(mapped.get(b"c".as_slice()), Some(&b"3".to_vec()));
2074 assert_eq!(mapped.get(b"d".as_slice()), Some(&b"4".to_vec()));
2075 assert!(!mapped.contains_key(b"b".as_slice()));
2076 });
2077 }
2078 #[test]
2079 fn async_default_ordered_batch_reads_deduplicate_duplicate_keys() {
2080 let store = DefaultAsyncReadStore::with_entries(1, &[(b"a", b"1"), (b"b", b"2")]);
2081 let keys: Vec<&[u8]> = vec![b"a", b"a", b"missing", b"missing", b"b"];
2082
2083 let values = block_on(store.batch_get_ordered(&keys)).unwrap();
2084
2085 assert_eq!(
2086 values,
2087 vec![
2088 Some(b"1".to_vec()),
2089 Some(b"1".to_vec()),
2090 None,
2091 None,
2092 Some(b"2".to_vec())
2093 ]
2094 );
2095 assert_eq!(
2096 store.get_calls.load(Ordering::Relaxed),
2097 3,
2098 "async ordered batch reads should point-read each unique key at most once"
2099 );
2100 }
2101 #[test]
2102 fn async_default_ordered_batch_reads_respect_read_parallelism() {
2103 let store = DefaultAsyncReadStore::with_entries(
2104 2,
2105 &[(b"a", b"1"), (b"b", b"2"), (b"c", b"3"), (b"d", b"4")],
2106 );
2107 let keys: Vec<&[u8]> = vec![b"a", b"b", b"c", b"d"];
2108
2109 let values = block_on(store.batch_get_ordered(&keys)).unwrap();
2110
2111 assert_eq!(
2112 values,
2113 vec![
2114 Some(b"1".to_vec()),
2115 Some(b"2".to_vec()),
2116 Some(b"3".to_vec()),
2117 Some(b"4".to_vec())
2118 ]
2119 );
2120 assert_eq!(store.get_calls.load(Ordering::Relaxed), 4);
2121 assert_eq!(
2122 store.max_in_flight.load(Ordering::Relaxed),
2123 2,
2124 "default async batch reads should cap concurrent point reads"
2125 );
2126 }
2127 #[test]
2128 fn arc_async_store_forwards_ordered_reads() {
2129 let store = std::sync::Arc::new(DefaultAsyncReadStore::with_entries(
2130 2,
2131 &[(b"a", b"1"), (b"b", b"2")],
2132 ));
2133 let keys: Vec<&[u8]> = vec![b"b", b"a"];
2134
2135 let values = block_on(store.batch_get_ordered(&keys)).unwrap();
2136
2137 assert_eq!(values, vec![Some(b"2".to_vec()), Some(b"1".to_vec())]);
2138 }
2139}