hashiverse_server_lib/environment/
mem_environment_store.rs

1//! # In-memory [`EnvironmentStore`] for tests
2//!
3//! Volatile implementation of
4//! [`crate::environment::environment_store::EnvironmentStore`] backed by a
5//! `HashMap` for bundles/metadata and a `BTreeMap` for feedback (so range queries
6//! over composite `(location, post, feedback_type)` keys work identically to the
7//! disk store).
8//!
9//! One subtlety: the decimation path needs to hold a write lock while *also*
10//! iterating the bundle set, and Rust's borrow rules don't allow that directly.
11//! The module uses an `ouroboros` self-referential struct to safely carry the
12//! guard alongside the iterator, matching what the disk store gets from its
13//! `fjall` cursor.
14
15use bytes::Bytes;
16use std::collections::{BTreeMap, HashMap};
17use std::sync::Arc;
18use async_trait::async_trait;
19use log::info;
20use ouroboros::self_referencing;
21use parking_lot::{RwLock,RwLockReadGuard};
22use hashiverse_lib::protocol::posting::encoded_post_feedback::EncodedPostFeedbackV1;
23use hashiverse_lib::tools::time::{TimeMillis, TimeMillisBytes};
24use hashiverse_lib::tools::types::Id;
25use crate::environment::environment::{Environment, EnvironmentDimensions, EnvironmentFactory, PostBundleMetadata};
26use crate::environment::environment_store::EnvironmentStore;
27
28pub struct MemEnvironmentFactory {
29    next_env_id: RwLock<usize>,
30}
31#[async_trait]
32impl EnvironmentFactory for MemEnvironmentFactory {
33    fn new(_base_path: &str) -> Self {
34        Self {
35            next_env_id: RwLock::new(0),
36        }
37    }
38    async fn open_next_available(&self, environment_dimensions: EnvironmentDimensions) -> anyhow::Result<Environment> {
39        let env_id = {
40            let mut next_env_id = self.next_env_id.write();
41            *next_env_id += 1;
42            *next_env_id
43        };
44
45        let mem_environment_store = Arc::new(MemEnvironmentStore::new(env_id));
46        let environment = Environment::new(mem_environment_store, environment_dimensions).await?;
47        Ok(environment)
48    }
49}
50
51/// MemEnvironmentStore implements the `EnvironmentStore` trait for use during testing.  Everything is stored in memory and is volatile to restarts.
52
53pub struct MemEnvironmentStore {
54    config: RwLock<HashMap<String, Vec<u8>>>,
55    post_bundle_metadata: RwLock<HashMap<Id, Vec<u8>>>,
56    post_bundle_last_accessed: RwLock<BTreeMap<Id, TimeMillisBytes>>,
57    post_bundle_bytes: RwLock<HashMap<Id, Bytes>>,
58    post_bundle_feedbacks: RwLock<HashMap<Id, Vec<EncodedPostFeedbackV1>>> // Map from post_bundle_location_id -> EncodedPostFeedbackV1
59}
60
61impl MemEnvironmentStore {
62    fn new(env_id: usize) -> Self {
63        info!("using MemEnvironmentStore {}", env_id);
64
65        Self {
66            config: RwLock::new(HashMap::new()),
67            post_bundle_metadata: RwLock::new(HashMap::new()),
68            post_bundle_last_accessed: RwLock::new(BTreeMap::new()),
69            post_bundle_bytes: RwLock::new(HashMap::new()),
70            post_bundle_feedbacks: RwLock::new(HashMap::new()),
71        }
72    }
73}
74
75impl EnvironmentStore for MemEnvironmentStore {
76
77    fn post_bundle_count(&self) -> anyhow::Result<usize> {
78        let len = self.post_bundle_last_accessed.read().len();
79        Ok(len)
80    }
81
82    fn post_bundle_metadata_get(&self, location_id: &Id) -> anyhow::Result<Option<PostBundleMetadata>> {
83        let post_bundle_metadata = self.post_bundle_metadata.read();
84        let bytes = post_bundle_metadata.get(location_id);
85        match bytes {
86            Some(bytes) => Ok(Some(postcard::from_bytes(bytes)?)),
87            None => Ok(None),
88        }
89    }
90
91    fn post_bundle_metadata_put(&self, location_id: &Id, post_bundle_metadata: &PostBundleMetadata) -> anyhow::Result<()> {
92        let bytes = postcard::to_stdvec(post_bundle_metadata)?;
93        self.post_bundle_metadata.write().insert(*location_id, bytes.to_vec());
94        Ok(())
95    }
96
97    fn post_bundle_bytes_get(&self, location_id: &Id) -> anyhow::Result<Option<Bytes>> {
98        Ok(self.post_bundle_bytes.read().get(location_id).cloned())
99    }
100
101    fn post_bundle_bytes_put(&self, location_id: &Id, bytes: &[u8]) -> anyhow::Result<()> {
102        self.post_bundle_bytes.write().insert(*location_id, Bytes::copy_from_slice(bytes));
103        Ok(())
104    }
105
106    fn post_bundles_last_accessed_flush(&self, post_bundles_last_accessed: &HashMap<Id, TimeMillis>) -> anyhow::Result<()> {
107        let mut self_post_bundle_last_accessed = self.post_bundle_last_accessed.write();
108        for (location_id, time_millis) in post_bundles_last_accessed.iter() {
109            self_post_bundle_last_accessed.insert(*location_id, time_millis.encode_be());
110        }
111        Ok(())
112    }
113
114    fn post_bundles_delete(&self, location_ids: &[Id]) -> anyhow::Result<()> {
115        let mut post_bundle_metadata = self.post_bundle_metadata.write();
116        let mut post_bundle_last_accessed = self.post_bundle_last_accessed.write();
117        let mut post_bundle_bytes = self.post_bundle_bytes.write();
118        let mut post_bundle_feedbacks = self.post_bundle_feedbacks.write();
119
120        for location_id in location_ids {
121            post_bundle_metadata.remove(location_id);
122            post_bundle_last_accessed.remove(location_id);
123            post_bundle_bytes.remove(location_id);
124            post_bundle_feedbacks.remove(location_id);
125        }
126
127        Ok(())
128    }
129
130    fn post_bundles_last_accessed_iter(&self, location_id: &Id) -> Box<dyn Iterator<Item = Result<(Id, TimeMillisBytes), anyhow::Error>> + '_> {
131        let post_bundle_last_accessed = self.post_bundle_last_accessed.read();
132
133        let iter = LockedPostBundleIterBuilder {
134            guard: post_bundle_last_accessed,
135            inner_builder: |post_bundle_last_accessed| {
136                let it = post_bundle_last_accessed.range(location_id..)
137                    .chain(post_bundle_last_accessed.range(..location_id))
138                    .map(|(k, v)| Ok((*k, *v)));
139                Box::new(it)
140            },
141        }.build();
142
143        Box::new(iter)
144    }
145
146    fn config_get_bytes(&self, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
147        Ok(self.config.read().get(key).map(|v| v.to_vec()))
148    }
149
150    fn config_put_bytes(&self, key: &str, v: Vec<u8>) -> anyhow::Result<()> {
151        self.config.write().insert(key.to_string(), v);
152        Ok(())
153    }
154
155    fn post_bundle_feedbacks_bytes_get(&self, post_bundle_location_id: &Id) -> anyhow::Result<Bytes> {
156        let mut bytes = Vec::new();
157        let feedbacks = self.post_bundle_feedbacks.read();
158        if let Some(entries) = feedbacks.get(post_bundle_location_id) {
159            for f in entries {
160                EncodedPostFeedbackV1::append_encode_direct_to_bytes(&mut bytes, f.post_id.as_ref(), f.feedback_type, f.salt.as_ref(), f.pow)?;
161            }
162        }
163        Ok(Bytes::from(bytes))
164    }
165
166    fn post_feedback_put_if_more_powerful(&self, location_id: &Id, encoded_post_feedback: &EncodedPostFeedbackV1) -> anyhow::Result<()> {
167        let mut feedbacks = self.post_bundle_feedbacks.write();
168        let entries = feedbacks.entry(*location_id).or_default();
169
170        // Find existing entry for this (post_id, feedback_type) pair
171        if let Some(existing) = entries.iter_mut().find(|f| f.post_id == encoded_post_feedback.post_id && f.feedback_type == encoded_post_feedback.feedback_type) {
172            if existing.pow >= encoded_post_feedback.pow {
173                // Already have equal or better feedback — do not downgrade
174                return Ok(());
175            }
176            *existing = encoded_post_feedback.clone();
177        } else {
178            entries.push(encoded_post_feedback.clone());
179        }
180
181        Ok(())
182    }
183}
184
185
186#[self_referencing]
187pub struct LockedPostBundleIter<'a> {
188    guard: RwLockReadGuard<'a, BTreeMap<Id, TimeMillisBytes>>,
189
190    #[borrows(guard)]
191    #[not_covariant] // Required for iterators
192    inner: Box<dyn Iterator<Item = Result<(Id, TimeMillisBytes), anyhow::Error>> + 'this>,
193}
194
195impl<'a> Iterator for LockedPostBundleIter<'a> {
196    type Item = Result<(Id, TimeMillisBytes), anyhow::Error>;
197
198    fn next(&mut self) -> Option<Self::Item> {
199        self.with_inner_mut(|it| it.next())
200    }
201}
202
203
204
205#[cfg(test)]
206mod tests {
207    use crate::environment;
208    use crate::environment::mem_environment_store::MemEnvironmentFactory;
209
210    #[tokio::test]
211    async fn basics_test() -> anyhow::Result<()> {
212        environment::environment::tests::basics_test::<MemEnvironmentFactory>().await
213    }
214
215    #[tokio::test]
216    async fn decimation_test() -> anyhow::Result<()> {
217        environment::environment::tests::decimation_test::<MemEnvironmentFactory>().await
218    }
219
220    #[tokio::test]
221    async fn decimation_convergence_test() -> anyhow::Result<()> {
222        environment::environment::tests::decimation_convergence_test::<MemEnvironmentFactory>(128 * 1000).await
223    }
224    #[tokio::test]
225    async fn decimation_existence_test() -> anyhow::Result<()> {
226        environment::environment::tests::decimation_existence_test::<MemEnvironmentFactory>().await
227    }
228    #[tokio::test]
229    async fn decimation_feedback_deleted_test() -> anyhow::Result<()> {
230        environment::environment::tests::decimation_feedback_deleted_test::<MemEnvironmentFactory>().await
231    }
232
233    #[tokio::test]
234    async fn feedback_bytes_get_test() -> anyhow::Result<()> {
235        environment::environment::tests::feedback_bytes_get_test::<MemEnvironmentFactory>().await
236    }
237
238    #[tokio::test]
239    async fn feedback_put_if_more_powerful_test() -> anyhow::Result<()> {
240        environment::environment::tests::feedback_put_if_more_powerful_test::<MemEnvironmentFactory>().await
241    }
242}
hashiverse_server_lib/environment/mem_environment_store.rs

hashiverse_server_lib/environment/
mem_environment_store.rs
