hashtree_network/
routed_store.rs1use std::sync::Arc;
4
5use async_trait::async_trait;
6use hashtree_core::store::StoreStats;
7use hashtree_core::{Hash, Store, StoreError};
8
9use crate::BlobRouter;
10
11pub struct RoutedStore<S: Store + ?Sized> {
18 primary: Arc<S>,
19 router: Arc<BlobRouter>,
20}
21
22impl<S: Store + ?Sized> RoutedStore<S> {
23 pub fn new(primary: Arc<S>, router: Arc<BlobRouter>) -> Self {
24 Self { primary, router }
25 }
26
27 pub fn primary(&self) -> &Arc<S> {
28 &self.primary
29 }
30
31 pub fn router(&self) -> &Arc<BlobRouter> {
32 &self.router
33 }
34}
35
36#[async_trait]
37impl<S: Store + ?Sized + 'static> Store for RoutedStore<S> {
38 async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
39 self.primary.put(hash, data).await
40 }
41
42 async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
43 self.primary.put_many(items).await
44 }
45
46 async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
47 self.router.get(hash, None).await
48 }
49
50 async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
51 Ok(self.get(hash).await?.is_some())
52 }
53
54 async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
55 self.primary.delete(hash).await
56 }
57
58 fn set_max_bytes(&self, max: u64) {
59 self.primary.set_max_bytes(max);
60 }
61
62 fn max_bytes(&self) -> Option<u64> {
63 self.primary.max_bytes()
64 }
65
66 async fn stats(&self) -> StoreStats {
67 self.primary.stats().await
68 }
69
70 async fn evict_if_needed(&self) -> Result<u64, StoreError> {
71 self.primary.evict_if_needed().await
72 }
73
74 async fn pin(&self, hash: &Hash) -> Result<(), StoreError> {
75 self.primary.pin(hash).await
76 }
77
78 async fn unpin(&self, hash: &Hash) -> Result<(), StoreError> {
79 self.primary.unpin(hash).await
80 }
81
82 fn pin_count(&self, hash: &Hash) -> u32 {
83 self.primary.pin_count(hash)
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use std::sync::atomic::{AtomicUsize, Ordering};
90 use std::time::Duration;
91
92 use super::*;
93 use crate::{BlobRouteEntry, BlobRouterConfig};
94 use hashtree_core::{sha256, BlobReply, BlobRequest, BlobRoute, MemoryStore};
95
96 struct DataRoute {
97 data: Vec<u8>,
98 calls: AtomicUsize,
99 }
100
101 #[async_trait]
102 impl BlobRoute for DataRoute {
103 async fn route(&self, _request: BlobRequest) -> Result<BlobReply, StoreError> {
104 self.calls.fetch_add(1, Ordering::SeqCst);
105 Ok(BlobReply::Data(self.data.clone()))
106 }
107 }
108
109 fn router(route: Arc<dyn BlobRoute>, cache: Arc<MemoryStore>) -> Arc<BlobRouter> {
110 Arc::new(
111 BlobRouter::new(
112 vec![BlobRouteEntry::new("remote", route)],
113 Some(cache),
114 BlobRouterConfig {
115 request_timeout: Duration::from_secs(1),
116 ..Default::default()
117 },
118 )
119 .unwrap(),
120 )
121 }
122
123 #[tokio::test]
124 async fn reads_use_the_router_and_its_verified_primary_cache() {
125 let data = b"routed store data".to_vec();
126 let hash = sha256(&data);
127 let primary = Arc::new(MemoryStore::new());
128 let remote = Arc::new(DataRoute {
129 data: data.clone(),
130 calls: AtomicUsize::new(0),
131 });
132 let store = RoutedStore::new(primary.clone(), router(remote.clone(), primary.clone()));
133
134 assert_eq!(store.get(&hash).await.unwrap(), Some(data.clone()));
135 assert_eq!(primary.get(&hash).await.unwrap(), Some(data));
136 assert_eq!(remote.calls.load(Ordering::SeqCst), 1);
137 }
138
139 #[tokio::test]
140 async fn mutations_only_touch_the_explicit_primary_store() {
141 let remote_data = b"remote".to_vec();
142 let remote = Arc::new(DataRoute {
143 data: remote_data.clone(),
144 calls: AtomicUsize::new(0),
145 });
146 let primary = Arc::new(MemoryStore::new());
147 let store = RoutedStore::new(primary.clone(), router(remote.clone(), primary.clone()));
148 let local_data = b"local write".to_vec();
149 let hash = sha256(&local_data);
150
151 assert!(store.put(hash, local_data.clone()).await.unwrap());
152 assert_eq!(primary.get(&hash).await.unwrap(), Some(local_data));
153 store.pin(&hash).await.unwrap();
154 assert_eq!(store.pin_count(&hash), 1);
155 assert!(store.delete(&hash).await.unwrap());
156 assert_eq!(remote.calls.load(Ordering::SeqCst), 0);
157 }
158}