cellular_raza_core/storage/
sled_database.rs

1use super::concepts::StorageError;
2use super::concepts::{StorageInterfaceLoad, StorageInterfaceOpen, StorageInterfaceStore};
3
4use serde::{Deserialize, Serialize};
5
6use std::collections::{BTreeMap, HashMap};
7use std::marker::PhantomData;
8
9/// Use the [sled] database to save results to an embedded database.
10// TODO use custom field for config [](https://docs.rs/sled/latest/sled/struct.Config.html) to let the user control these parameters
11#[derive(Clone)]
12pub struct SledStorageInterface<Id, Element, const TEMP: bool = false> {
13    db: sled::Db,
14    // TODO use this buffer
15    // buffer: StorageBuffer<Id, Element>,
16    id_phantom: PhantomData<Id>,
17    element_phantom: PhantomData<Element>,
18    bincode_config: bincode::config::Configuration,
19}
20
21impl<Id: core::fmt::Debug, Element: core::fmt::Debug, const TEMP: bool> core::fmt::Debug
22    for SledStorageInterface<Id, Element, TEMP>
23{
24    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
25        #[allow(unused)]
26        #[derive(Debug)]
27        struct DbDebug<Id, Element, const TEMP: bool> {
28            db: sled::Db,
29            id_phantom: PhantomData<Id>,
30            element_phantom: PhantomData<Element>,
31        }
32        let SledStorageInterface {
33            db,
34            id_phantom,
35            element_phantom,
36            bincode_config: _,
37        } = self;
38        let db_debug = DbDebug::<Id, Element, TEMP> {
39            db: db.clone(),
40            id_phantom: *id_phantom,
41            element_phantom: *element_phantom,
42        };
43        core::fmt::Debug::fmt(&db_debug, f)
44    }
45}
46
47impl<Id, Element, const TEMP: bool> SledStorageInterface<Id, Element, TEMP> {
48    /// Transform a u64 value to an iteration key which can be given to a sled tree.
49    fn iteration_to_key(iteration: u64) -> [u8; 8] {
50        iteration.to_le_bytes()
51    }
52
53    /// Transform the key given by the tree to the corresponding iteartion u64 value
54    fn key_to_iteration(
55        key: &sled::IVec,
56        bincode_config: bincode::config::Configuration,
57    ) -> Result<u64, StorageError> {
58        let iteration = bincode::serde::decode_from_slice(key, bincode_config)?.0;
59        Ok(iteration)
60    }
61
62    /// Get the correct tree of the iteration or create if not currently present.
63    fn open_or_create_tree(&self, iteration: u64) -> Result<sled::Tree, StorageError> {
64        let tree_key = Self::iteration_to_key(iteration);
65        let tree = self.db.open_tree(&tree_key)?;
66        Ok(tree)
67    }
68
69    fn open_tree(&self, iteration: u64) -> Result<Option<sled::Tree>, StorageError> {
70        let tree_key = Self::iteration_to_key(iteration);
71        if !self.db.tree_names().contains(&sled::IVec::from(&tree_key)) {
72            Ok(None)
73        } else {
74            let tree = self.db.open_tree(tree_key)?;
75            Ok(Some(tree))
76        }
77    }
78}
79
80impl<Id, Element, const TEMP: bool> StorageInterfaceOpen
81    for SledStorageInterface<Id, Element, TEMP>
82{
83    fn open_or_create(
84        location: &std::path::Path,
85        _storage_instance: u64,
86    ) -> Result<Self, StorageError> {
87        let config = sled::Config::default()
88            .mode(sled::Mode::HighThroughput)
89            .cache_capacity(1024 * 1024 * 1024 * 5) // 5gb
90            .path(&location)
91            .temporary(TEMP)
92            .use_compression(false);
93
94        let db = config.open()?;
95        let bincode_config = bincode::config::standard();
96
97        Ok(SledStorageInterface {
98            db,
99            id_phantom: PhantomData,
100            element_phantom: PhantomData,
101            bincode_config,
102        })
103    }
104
105    fn clone_to_new_instance(&self, _storage_instance: u64) -> Self {
106        Self {
107            db: self.db.clone(),
108            bincode_config: self.bincode_config.clone(),
109            id_phantom: PhantomData,
110            element_phantom: PhantomData,
111        }
112    }
113}
114
115impl<Id, Element, const TEMP: bool> StorageInterfaceStore<Id, Element>
116    for SledStorageInterface<Id, Element, TEMP>
117{
118    fn store_single_element(
119        &mut self,
120        iteration: u64,
121        identifier: &Id,
122        element: &Element,
123    ) -> Result<(), StorageError>
124    where
125        Id: Serialize,
126        Element: Serialize,
127    {
128        let tree = self.open_or_create_tree(iteration)?;
129
130        // Serialize the identifier and the element
131        let identifier_serialized =
132            bincode::serde::encode_to_vec(&identifier, self.bincode_config)?;
133        let element_serialized = bincode::serde::encode_to_vec(&element, self.bincode_config)?;
134        match tree.insert(identifier_serialized, element_serialized)? {
135            None => Ok(()),
136            Some(_) => Err(StorageError::InitError(format!(
137                "Element already present at iteration {}",
138                iteration
139            ))),
140        }?;
141        Ok(())
142    }
143
144    fn store_batch_elements<'a, I>(
145        &'a mut self,
146        iteration: u64,
147        identifiers_elements: I,
148    ) -> Result<(), StorageError>
149    where
150        Id: 'a + Serialize,
151        Element: 'a + Serialize,
152        I: Clone + IntoIterator<Item = (&'a Id, &'a Element)>,
153    {
154        let tree = self.open_or_create_tree(iteration)?;
155        let mut batch = sled::Batch::default();
156        for (identifier, element) in identifiers_elements.into_iter() {
157            let identifier_serialized =
158                bincode::serde::encode_to_vec(&identifier, self.bincode_config)?;
159            let element_serialized = bincode::serde::encode_to_vec(&element, self.bincode_config)?;
160            batch.insert(identifier_serialized, element_serialized)
161        }
162        tree.apply_batch(batch)?;
163        Ok(())
164    }
165}
166
167impl<Id, Element, const TEMP: bool> StorageInterfaceLoad<Id, Element>
168    for SledStorageInterface<Id, Element, TEMP>
169{
170    fn load_single_element(
171        &self,
172        iteration: u64,
173        identifier: &Id,
174    ) -> Result<Option<Element>, StorageError>
175    where
176        Id: Serialize + for<'a> Deserialize<'a>,
177        Element: for<'a> Deserialize<'a>,
178    {
179        let tree = match self.open_tree(iteration)? {
180            Some(tree) => tree,
181            None => return Ok(None),
182        };
183        let identifier_serialized = bincode::serde::encode_to_vec(identifier, self.bincode_config)?;
184        match tree.get(&identifier_serialized)? {
185            Some(element_serialized) => {
186                let element: Element =
187                    bincode::serde::decode_from_slice(&element_serialized, self.bincode_config)?.0;
188                Ok(Some(element))
189            }
190            None => Ok(None),
191        }
192    }
193
194    fn load_element_history(&self, identifier: &Id) -> Result<HashMap<u64, Element>, StorageError>
195    where
196        Id: Serialize,
197        Element: for<'a> Deserialize<'a>,
198    {
199        // Keep track if the element has not been found in a database.
200        // If so we can either get the current minimal iteration or maximal depending on where it was found else.
201        let mut minimal_iteration = None;
202        let mut maximal_iteration = None;
203        let mut success_iteration = None;
204
205        // Save results in this hashmap
206        let mut accumulator = HashMap::new();
207        // Serialize the identifier
208        let identifier_serialized = bincode::serde::encode_to_vec(identifier, self.bincode_config)?;
209        for iteration_serialized in self.db.tree_names() {
210            // If we are above the maximal or below the minimal iteration, we skip checking
211            let iteration: u64 =
212                bincode::serde::decode_from_slice(&iteration_serialized, self.bincode_config)?.0;
213            match minimal_iteration {
214                None => (),
215                Some(min_iter) => {
216                    if iteration < min_iter {
217                        continue;
218                    }
219                }
220            }
221            match maximal_iteration {
222                None => (),
223                Some(max_iter) => {
224                    if max_iter < iteration {
225                        continue;
226                    }
227                }
228            }
229            // Get the tree for a random iteration
230            let tree = self.db.open_tree(iteration_serialized)?;
231            match tree.get(&identifier_serialized)? {
232                // We found and element insert it
233                Some(element_serialized) => {
234                    let element: Element = bincode::serde::decode_from_slice(
235                        &element_serialized,
236                        self.bincode_config,
237                    )?
238                    .0;
239                    accumulator.insert(iteration, element);
240                    success_iteration = Some(iteration);
241                }
242                // We did not find an element. Thus update the helper variables atop.
243                None => match (minimal_iteration, maximal_iteration, success_iteration) {
244                    (None, None, Some(suc_iter)) => {
245                        if iteration > suc_iter {
246                            maximal_iteration = Some(iteration);
247                        }
248                        if iteration < suc_iter {
249                            minimal_iteration = Some(iteration);
250                        }
251                    }
252                    (Some(min_iter), None, Some(suc_iter)) => {
253                        if iteration > suc_iter {
254                            maximal_iteration = Some(iteration);
255                        }
256                        if iteration < suc_iter && iteration > min_iter {
257                            minimal_iteration = Some(iteration);
258                        }
259                    }
260                    (None, Some(max_iter), Some(suc_iter)) => {
261                        if iteration > suc_iter && iteration < max_iter {
262                            maximal_iteration = Some(iteration);
263                        }
264                        if iteration < suc_iter {
265                            minimal_iteration = Some(iteration);
266                        }
267                    }
268                    (Some(min_iter), Some(max_iter), Some(suc_iter)) => {
269                        if iteration > suc_iter && iteration < max_iter {
270                            maximal_iteration = Some(iteration);
271                        }
272                        if iteration < suc_iter && iteration > min_iter {
273                            minimal_iteration = Some(iteration);
274                        }
275                    }
276                    (_, _, None) => (),
277                },
278            };
279        }
280        Ok(accumulator)
281    }
282
283    fn load_all_elements_at_iteration(
284        &self,
285        iteration: u64,
286    ) -> Result<HashMap<Id, Element>, StorageError>
287    where
288        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
289        Element: for<'a> Deserialize<'a>,
290    {
291        let tree = match self.open_tree(iteration)? {
292            Some(tree) => tree,
293            None => return Ok(HashMap::new()),
294        };
295        tree.iter()
296            .map(|entry_result| {
297                let (identifier_serialized, element_serialized) = entry_result?;
298                let identifier: Id =
299                    bincode::serde::decode_from_slice(&identifier_serialized, self.bincode_config)?
300                        .0;
301                let element: Element =
302                    bincode::serde::decode_from_slice(&element_serialized, self.bincode_config)?.0;
303                Ok((identifier, element))
304            })
305            .collect::<Result<HashMap<Id, Element>, StorageError>>()
306    }
307
308    fn load_all_elements(&self) -> Result<BTreeMap<u64, HashMap<Id, Element>>, StorageError>
309    where
310        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
311        Element: for<'a> Deserialize<'a>,
312    {
313        self.db
314            .tree_names()
315            .iter()
316            .map(|tree_name_serialized| {
317                let tree = self.db.open_tree(tree_name_serialized)?;
318                let iteration = Self::key_to_iteration(tree_name_serialized, self.bincode_config)?;
319                let identifier_to_element = tree
320                    .iter()
321                    .map(|entry_result| {
322                        let (identifier_serialized, element_serialized) = entry_result?;
323                        let identifier: Id = bincode::serde::decode_from_slice(
324                            &identifier_serialized,
325                            self.bincode_config,
326                        )?
327                        .0;
328                        let element: Element = bincode::serde::decode_from_slice(
329                            &element_serialized,
330                            self.bincode_config,
331                        )?
332                        .0;
333                        Ok((identifier, element))
334                    })
335                    .collect::<Result<HashMap<Id, Element>, StorageError>>()?;
336                Ok((iteration, identifier_to_element))
337            })
338            .collect::<Result<BTreeMap<u64, HashMap<Id, Element>>, StorageError>>()
339    }
340
341    fn get_all_iterations(&self) -> Result<Vec<u64>, StorageError> {
342        let iterations = self
343            .db
344            .tree_names()
345            .iter()
346            // TODO this should not be here! Fix it properly (I asked on sled discord)
347            .filter(|key| {
348                **key
349                    != sled::IVec::from(&[
350                        95, 95, 115, 108, 101, 100, 95, 95, 100, 101, 102, 97, 117, 108, 116,
351                    ])
352            })
353            .map(|tree_name_serialized| Self::key_to_iteration(tree_name_serialized, self.bincode_config))
354            .collect::<Result<Vec<_>, StorageError>>()?;
355
356        Ok(iterations)
357    }
358}
cellular_raza_core/storage/sled_database.rs

cellular_raza_core/storage/
sled_database.rs
