use core::{marker::PhantomData, fmt::Debug}; use std::{ sync::{Arc, RwLock}, collections::HashMap, }; use crate::{Plan, coins::Coin}; pub trait DbTxn: Send + Sync + Clone + Debug { fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>); fn get(&self, key: impl AsRef<[u8]>) -> Option>; fn del(&mut self, key: impl AsRef<[u8]>); fn commit(self); } pub trait Db: 'static + Send + Sync + Clone + Debug { type Transaction: DbTxn; fn key(db_dst: &'static [u8], item_dst: &'static [u8], key: impl AsRef<[u8]>) -> Vec { let db_len = u8::try_from(db_dst.len()).unwrap(); let dst_len = u8::try_from(item_dst.len()).unwrap(); [[db_len].as_ref(), db_dst, [dst_len].as_ref(), item_dst, key.as_ref()].concat().to_vec() } fn txn(&mut self) -> Self::Transaction; fn get(&self, key: impl AsRef<[u8]>) -> Option>; } // TODO: Replace this with RocksDB #[derive(Clone, Debug)] pub struct MemDb(Arc, Vec>>>); impl MemDb { #[allow(clippy::new_without_default)] pub fn new() -> MemDb { MemDb(Arc::new(RwLock::new(HashMap::new()))) } } impl DbTxn for MemDb { fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) { self.0.write().unwrap().insert(key.as_ref().to_vec(), value.as_ref().to_vec()); } fn get(&self, key: impl AsRef<[u8]>) -> Option> { self.0.read().unwrap().get(key.as_ref()).cloned() } fn del(&mut self, key: impl AsRef<[u8]>) { self.0.write().unwrap().remove(key.as_ref()); } fn commit(self) {} } impl Db for MemDb { type Transaction = MemDb; fn txn(&mut self) -> MemDb { Self(self.0.clone()) } fn get(&self, key: impl AsRef<[u8]>) -> Option> { self.0.read().unwrap().get(key.as_ref()).cloned() } } #[derive(Debug)] pub struct MainDb(D, PhantomData); impl MainDb { pub fn new(db: D) -> Self { Self(db, PhantomData) } fn main_key(dst: &'static [u8], key: impl AsRef<[u8]>) -> Vec { D::key(b"MAIN", dst, key) } fn plan_key(id: &[u8]) -> Vec { Self::main_key(b"plan", id) } fn signing_key(key: &[u8]) -> Vec { Self::main_key(b"signing", key) } pub fn save_signing(&mut self, key: &[u8], block_number: u64, time: u64, plan: &Plan) { let id = plan.id(); // Creating a TXN here is arguably an anti-pattern, yet nothing here expects atomicity let mut txn = self.0.txn(); { let mut signing = txn.get(Self::signing_key(key)).unwrap_or(vec![]); // If we've already noted we're signing this, return assert_eq!(signing.len() % 32, 0); for i in 0 .. (signing.len() / 32) { if signing[(i * 32) .. ((i + 1) * 32)] == id { return; } } signing.extend(&id); txn.put(Self::signing_key(key), id); } { let mut buf = block_number.to_le_bytes().to_vec(); buf.extend(&time.to_le_bytes()); plan.write(&mut buf).unwrap(); txn.put(Self::plan_key(&id), &buf); } txn.commit(); } pub fn signing(&self, key: &[u8]) -> Vec<(u64, u64, Plan)> { let signing = self.0.get(Self::signing_key(key)).unwrap_or(vec![]); let mut res = vec![]; assert_eq!(signing.len() % 32, 0); for i in 0 .. (signing.len() / 32) { let id = &signing[(i * 32) .. ((i + 1) * 32)]; let buf = self.0.get(Self::plan_key(id)).unwrap(); let block_number = u64::from_le_bytes(buf[.. 8].try_into().unwrap()); let time = u64::from_le_bytes(buf[8 .. 16].try_into().unwrap()); let plan = Plan::::read::<&[u8]>(&mut &buf[16 ..]).unwrap(); assert_eq!(id, &plan.id()); res.push((block_number, time, plan)); } res } pub fn finish_signing(&mut self, key: &[u8], id: [u8; 32]) { let mut signing = self.0.get(Self::signing_key(key)).unwrap_or(vec![]); assert_eq!(signing.len() % 32, 0); let mut found = false; for i in 0 .. (signing.len() / 32) { let start = i * 32; let end = i + 32; if signing[start .. end] == id { found = true; signing = [&signing[.. start], &signing[end ..]].concat().to_vec(); break; } } if !found { log::warn!("told to finish signing {} yet wasn't actively signing it", hex::encode(id)); } let mut txn = self.0.txn(); txn.put(Self::signing_key(key), signing); txn.commit(); } }