balius_runtime/
store.rs

1use itertools::Itertools;
2use prost::Message;
3use redb::{ReadableTable as _, TableDefinition, WriteTransaction};
4use std::{path::Path, sync::Arc};
5use tracing::warn;
6
7use crate::{Block, ChainPoint, Error};
8
9pub type WorkerId = String;
10pub type LogSeq = u64;
11
12#[derive(Message)]
13pub struct LogEntry {
14    #[prost(bytes, tag = "1")]
15    pub next_block: Vec<u8>,
16    #[prost(bytes, repeated, tag = "2")]
17    pub undo_blocks: Vec<Vec<u8>>,
18}
19
20impl redb::Value for LogEntry {
21    type SelfType<'a>
22        = LogEntry
23    where
24        Self: 'a;
25
26    type AsBytes<'a>
27        = Vec<u8>
28    where
29        Self: 'a;
30
31    fn fixed_width() -> Option<usize> {
32        None
33    }
34
35    fn from_bytes<'a>(data: &'a [u8]) -> Self::SelfType<'a>
36    where
37        Self: 'a,
38    {
39        prost::Message::decode(data).unwrap()
40    }
41
42    fn as_bytes<'a, 'b: 'a>(value: &'a Self::SelfType<'b>) -> Self::AsBytes<'a>
43    where
44        Self: 'a,
45        Self: 'b,
46    {
47        value.encode_to_vec()
48    }
49
50    fn type_name() -> redb::TypeName {
51        redb::TypeName::new("LogEntry")
52    }
53}
54
55const CURSORS: TableDefinition<WorkerId, LogSeq> = TableDefinition::new("cursors");
56const WAL: TableDefinition<LogSeq, LogEntry> = TableDefinition::new("wal");
57
58const DEFAULT_CACHE_SIZE_MB: usize = 50;
59
60pub struct AtomicUpdate {
61    wx: WriteTransaction,
62    log_seq: LogSeq,
63}
64
65impl AtomicUpdate {
66    pub fn update_worker_cursor(&mut self, id: &str) -> Result<(), super::Error> {
67        let mut table = self.wx.open_table(CURSORS)?;
68        table.insert(id.to_owned(), self.log_seq)?;
69
70        Ok(())
71    }
72
73    pub fn commit(self) -> Result<(), super::Error> {
74        self.wx.commit()?;
75        Ok(())
76    }
77}
78
79#[derive(Clone)]
80pub struct Store {
81    db: Arc<redb::Database>,
82    log_seq: LogSeq,
83}
84
85impl Store {
86    pub fn open(path: impl AsRef<Path>, cache_size: Option<usize>) -> Result<Self, super::Error> {
87        let inner = redb::Database::builder()
88            .set_repair_callback(|x| {
89                warn!(progress = x.progress() * 100f64, "balius db is repairing")
90            })
91            .set_cache_size(1024 * 1024 * cache_size.unwrap_or(DEFAULT_CACHE_SIZE_MB))
92            .create(path)?;
93
94        let log_seq = Self::load_log_seq(&inner)?.unwrap_or_default();
95
96        let out = Self {
97            db: Arc::new(inner),
98            log_seq,
99        };
100
101        Ok(out)
102    }
103
104    fn load_log_seq(db: &redb::Database) -> Result<Option<LogSeq>, Error> {
105        let rx = db.begin_read()?;
106
107        match rx.open_table(WAL) {
108            Ok(table) => {
109                let last = table.last()?;
110                Ok(last.map(|(k, _)| k.value()))
111            }
112            Err(redb::TableError::TableDoesNotExist(_)) => Ok(None),
113            Err(e) => return Err(e.into()),
114        }
115    }
116
117    fn get_entry(&self, seq: LogSeq) -> Result<Option<LogEntry>, Error> {
118        let rx = self.db.begin_read()?;
119        let table = rx.open_table(WAL)?;
120        let entry = table.get(seq)?;
121        Ok(entry.map(|x| x.value()))
122    }
123
124    pub fn find_chain_point(&self, seq: LogSeq) -> Result<Option<ChainPoint>, Error> {
125        let entry = self.get_entry(seq)?;
126        let block = Block::from_bytes(&entry.unwrap().next_block);
127
128        Ok(Some(block.chain_point()))
129    }
130
131    pub fn write_ahead(
132        &mut self,
133        undo_blocks: &Vec<Block>,
134        next_block: &Block,
135    ) -> Result<LogSeq, Error> {
136        self.log_seq += 1;
137
138        let wx = self.db.begin_write()?;
139        {
140            wx.open_table(WAL)?.insert(
141                self.log_seq,
142                LogEntry {
143                    next_block: next_block.to_bytes(),
144                    undo_blocks: undo_blocks.iter().map(|x| x.to_bytes()).collect(),
145                },
146            )?;
147        }
148
149        wx.commit()?;
150        Ok(self.log_seq)
151    }
152
153    // TODO: see if loading in batch is worth it
154    pub fn get_worker_cursor(&self, id: &str) -> Result<Option<LogSeq>, super::Error> {
155        let rx = self.db.begin_read()?;
156
157        let table = match rx.open_table(CURSORS) {
158            Ok(table) => table,
159            Err(redb::TableError::TableDoesNotExist(_)) => return Ok(None),
160            Err(e) => return Err(e.into()),
161        };
162
163        let cursor = table.get(id.to_owned())?;
164        Ok(cursor.map(|x| x.value()))
165    }
166
167    pub fn start_atomic_update(&self, log_seq: LogSeq) -> Result<AtomicUpdate, super::Error> {
168        let wx = self.db.begin_write()?;
169        Ok(AtomicUpdate { wx, log_seq })
170    }
171
172    // TODO: I don't think we need this since we're going to load each cursor as
173    // part of the loaded worker
174    pub fn lowest_cursor(&self) -> Result<Option<LogSeq>, super::Error> {
175        let rx = self.db.begin_read()?;
176
177        let table = rx.open_table(CURSORS)?;
178
179        let cursors: Vec<_> = table
180            .iter()?
181            .map_ok(|(_, value)| value.value())
182            .try_collect()?;
183
184        let lowest = cursors.iter().fold(None, |all, item| all.min(Some(*item)));
185
186        Ok(lowest)
187    }
188}
balius_runtime/store.rs

balius_runtime/
store.rs
