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
15#[async_trait::async_trait]
16impl KvProvider for MemoryKv {
17    async fn get_value(&mut self, worker_id: &str, key: String) -> Result<Payload, KvError> {
18        match self.map.entry(worker_id.to_string()).or_default().get(&key) {
19            Some(value) => Ok(value.clone()),
20            None => Err(KvError::NotFound(key)),
21        }
22    }
23
24    async fn set_value(
25        &mut self,
26        worker_id: &str,
27        key: String,
28        value: Payload,
29    ) -> Result<(), KvError> {
30        let _ = self
31            .map
32            .entry(worker_id.to_string())
33            .or_default()
34            .insert(key, value);
35        Ok(())
36    }
37
38    async fn list_values(
39        &mut self,
40        worker_id: &str,
41        prefix: String,
42    ) -> Result<Vec<String>, KvError> {
43        let range = self
44            .map
45            .entry(worker_id.to_string())
46            .or_default()
47            .range((Bound::Included(prefix.clone()), Bound::Unbounded));
48        let mut result = vec![];
49        for (k, _) in range {
50            if k.starts_with(&prefix) {
51                result.push(k.to_string());
52            } else {
53                // Sorted, if prefix doesn't match then we break
54                break;
55            }
56        }
57        Ok(result)
58    }
59}
balius_runtime/kv/memory.rs

balius_runtime/kv/
memory.rs
