balius_runtime/kv/
memory.rs

1/// In memory implementation of KV for development usage.
2use std::collections::BTreeMap;
3
4use crate::wit::balius::app::kv as wit;
5use std::ops::Bound;
6use wit::{KvError, Payload};
7
8use super::KvProvider;
9
10#[derive(Default, Clone)]
11pub struct MemoryKv {
12    map: BTreeMap<String, BTreeMap<String, Vec<u8>>>,
13}
14
15impl From<BTreeMap<String, BTreeMap<String, Vec<u8>>>> for MemoryKv {
16    fn from(value: BTreeMap<String, BTreeMap<String, Vec<u8>>>) -> Self {
17        Self { map: value }
18    }
19}
20
21#[async_trait::async_trait]
22impl KvProvider for MemoryKv {
23    async fn get_value(&mut self, worker_id: &str, key: String) -> Result<Payload, KvError> {
24        match self.map.entry(worker_id.to_string()).or_default().get(&key) {
25            Some(value) => Ok(value.clone()),
26            None => Err(KvError::NotFound(key)),
27        }
28    }
29
30    async fn set_value(
31        &mut self,
32        worker_id: &str,
33        key: String,
34        value: Payload,
35    ) -> Result<(), KvError> {
36        let _ = self
37            .map
38            .entry(worker_id.to_string())
39            .or_default()
40            .insert(key, value);
41        Ok(())
42    }
43
44    async fn list_values(
45        &mut self,
46        worker_id: &str,
47        prefix: String,
48    ) -> Result<Vec<String>, KvError> {
49        let range = self
50            .map
51            .entry(worker_id.to_string())
52            .or_default()
53            .range((Bound::Included(prefix.clone()), Bound::Unbounded));
54        let mut result = vec![];
55        for (k, _) in range {
56            if k.starts_with(&prefix) {
57                result.push(k.to_string());
58            } else {
59                // Sorted, if prefix doesn't match then we break
60                break;
61            }
62        }
63        Ok(result)
64    }
65}
balius_runtime/kv/memory.rs

balius_runtime/kv/
memory.rs
