From 114c879c6ff19252480528fd5de007070f5fbb91 Mon Sep 17 00:00:00 2001 From: projectmoon Date: Tue, 20 Oct 2020 20:54:20 +0000 Subject: [PATCH] Count user variables on a per-room basis. --- src/db.rs | 272 ++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 256 insertions(+), 16 deletions(-) diff --git a/src/db.rs b/src/db.rs index 4af7474..46c463b 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,5 +1,7 @@ use byteorder::LittleEndian; -use sled::{Db, IVec, Tree}; +use sled::transaction::abort; +use sled::transaction::{TransactionError, TransactionalTree, UnabortableTransactionError}; +use sled::{Db, Tree}; use std::collections::HashMap; use thiserror::Error; use zerocopy::byteorder::I32; @@ -10,6 +12,8 @@ use zerocopy::{AsBytes, LayoutVerified}; /// read. type LittleEndianI32Layout<'a> = LayoutVerified<&'a [u8], I32>; +const VARIABLE_COUNT_KEY: &'static str = "variable_count"; + #[derive(Clone)] pub struct Database { db: Db, @@ -32,6 +36,23 @@ pub enum DataError { #[error("internal database error: {0}")] InternalError(#[from] sled::Error), + + #[error("transaction error: {0}")] + TransactionError(#[from] sled::transaction::TransactionError), + + #[error("unabortable transaction error: {0}")] + UnabortableTransactionError(#[from] UnabortableTransactionError), +} + +impl From> for DataError { + fn from(error: TransactionError) -> Self { + match error { + TransactionError::Abort(data_err) => data_err, + TransactionError::Storage(storage_err) => { + DataError::TransactionError(TransactionError::Storage(storage_err)) + } + } + } } fn to_key(room_id: &str, username: &str, variable_name: &str) -> Vec { @@ -42,6 +63,10 @@ fn to_key(room_id: &str, username: &str, variable_name: &str) -> Vec { key } +fn room_variable_count_key(room_id: &str, username: &str) -> Vec { + to_key(room_id, username, VARIABLE_COUNT_KEY) +} + fn to_prefix(room_id: &str, username: &str) -> Vec { let mut prefix = vec![]; prefix.extend_from_slice(room_id.as_bytes()); @@ -49,7 +74,9 @@ fn to_prefix(room_id: &str, username: &str) -> Vec { prefix } -fn convert(raw_value: &[u8]) -> Result { +/// Convert bytes to an i32 with zero-copy deserialization. An error +/// is returned if the bytes do not represent an i32. +fn convert_i32(raw_value: &[u8]) -> Result { let layout = LittleEndianI32Layout::new_unaligned(raw_value.as_ref()); if let Some(layout) = layout { @@ -60,6 +87,29 @@ fn convert(raw_value: &[u8]) -> Result { } } +/// Atomically alter the count of variables in the database, by the +/// given amount. Count cannot go below 0. +fn alter_room_variable_count( + variables: &TransactionalTree, + room_id: &str, + username: &str, + amount: i32, +) -> Result { + let key = room_variable_count_key(room_id, username); + let mut new_count = match variables.get(&key)? { + Some(bytes) => convert_i32(&bytes)? + amount, + None => amount, + }; + + if new_count < 0 { + new_count = 0; + } + + let db_value: I32 = I32::new(new_count); + variables.insert(key, db_value.as_bytes())?; + Ok(new_count) +} + impl Database { pub fn new>(path: P) -> Result { let db = sled::open(path)?; @@ -89,17 +139,30 @@ impl Database { //Strips room and username from key, leaving //behind name. let variable_name = std::str::from_utf8(&key[prefix_len..])?; - Ok((variable_name.to_owned(), convert(&raw_value)?)) + Ok((variable_name.to_owned(), convert_i32(&raw_value)?)) } Err(e) => Err(e.into()), }) .collect(); - //Convert to hash map. Can we do this in the first mapping - //step instead? + //Convert_I32 to hash map. Can we do this in the first mapping + //step instead? For some reason this is faster. variables.map(|entries| entries.into_iter().collect()) } + pub async fn get_variable_count( + &self, + room_id: &str, + username: &str, + ) -> Result { + let key = room_variable_count_key(room_id, username); + if let Some(raw_value) = self.variables.get(&key)? { + convert_i32(&raw_value) + } else { + Ok(0) + } + } + pub async fn get_user_variable( &self, room_id: &str, @@ -109,7 +172,7 @@ impl Database { let key = to_key(room_id, username, variable_name); if let Some(raw_value) = self.variables.get(&key)? { - convert(&raw_value) + convert_i32(&raw_value) } else { Err(DataError::KeyDoesNotExist(String::from_utf8(key).unwrap())) } @@ -122,11 +185,23 @@ impl Database { variable_name: &str, value: i32, ) -> Result<(), DataError> { - let key = to_key(room_id, username, variable_name); - let db_value: I32 = I32::new(value); self.variables - .insert(&key, IVec::from(db_value.as_bytes()))?; - Ok(()) + .transaction(|tx| { + let key = to_key(room_id, username, variable_name); + let db_value: I32 = I32::new(value); + let old_value = tx.insert(key, db_value.as_bytes())?; + + //Only increment variable count on new keys. + if let None = old_value { + match alter_room_variable_count(&tx, room_id, username, 1) { + Err(e) => abort(e), + _ => Ok(()), + } + } else { + Ok(()) + } + }) + .map_err(|e| e.into()) } pub async fn delete_user_variable( @@ -135,11 +210,176 @@ impl Database { username: &str, variable_name: &str, ) -> Result<(), DataError> { - let key = to_key(room_id, username, variable_name); - if let Some(_) = self.variables.remove(&key)? { - Ok(()) - } else { - Err(DataError::KeyDoesNotExist(String::from_utf8(key).unwrap())) - } + self.variables + .transaction(|tx| { + let key = to_key(room_id, username, variable_name); + if let Some(_) = tx.remove(key.clone())? { + match alter_room_variable_count(&tx, room_id, username, -1) { + Err(e) => abort(e), + _ => Ok(()), + } + } else { + abort(DataError::KeyDoesNotExist(String::from_utf8(key).unwrap())) + } + }) + .map_err(|e| e.into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + //Room Variable count tests + + #[tokio::test] + async fn alter_room_variable_count_test() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + + let alter_count = |amount: i32| { + db.variables + .transaction(|tx| { + match alter_room_variable_count(&tx, "room", "username", amount) { + Err(e) => abort(e), + _ => Ok(()), + } + }) + .expect("got transaction failure"); + }; + + async fn get_count(db: &Database) -> i32 { + db.get_variable_count("room", "username") + .await + .expect("could not get variable count") + } + + //addition + alter_count(5); + assert_eq!(5, get_count(&db).await); + + //subtraction + alter_count(-3); + assert_eq!(2, get_count(&db).await); + } + + #[tokio::test] + async fn alter_room_variable_count_cannot_go_below_0_test() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + db.variables + .transaction( + |tx| match alter_room_variable_count(&tx, "room", "username", -1000) { + Err(e) => abort(e), + _ => Ok(()), + }, + ) + .expect("got transaction failure"); + + let count = db + .get_variable_count("room", "username") + .await + .expect("could not get variable count"); + + assert_eq!(0, count); + } + + #[tokio::test] + async fn empty_db_reports_0_room_variable_count_test() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + + let count = db + .get_variable_count("room", "username") + .await + .expect("could not get variable count"); + + assert_eq!(0, count); + } + + #[tokio::test] + async fn set_user_variable_increments_count() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + db.set_user_variable("room", "username", "myvariable", 5) + .await + .expect("could not insert variable"); + + let count = db + .get_variable_count("room", "username") + .await + .expect("could not get variable count"); + + assert_eq!(1, count); + } + + #[tokio::test] + async fn update_user_variable_does_not_increment_count() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + db.set_user_variable("room", "username", "myvariable", 5) + .await + .expect("could not insert variable"); + + db.set_user_variable("room", "username", "myvariable", 10) + .await + .expect("could not update variable"); + + let count = db + .get_variable_count("room", "username") + .await + .expect("could not get variable count"); + + assert_eq!(1, count); + } + + // Set/get/delete variable tests + + #[tokio::test] + async fn set_and_get_variable_test() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + db.set_user_variable("room", "username", "myvariable", 5) + .await + .expect("could not insert variable"); + + let value = db + .get_user_variable("room", "username", "myvariable") + .await + .expect("could not get value"); + + assert_eq!(5, value); + } + + #[tokio::test] + async fn delete_variable_test() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + db.set_user_variable("room", "username", "myvariable", 5) + .await + .expect("could not insert variable"); + + db.delete_user_variable("room", "username", "myvariable") + .await + .expect("could not delete value"); + + let result = db.get_user_variable("room", "username", "myvariable").await; + + assert!(result.is_err()); + assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); + } + + #[tokio::test] + async fn get_missing_variable_returns_key_does_not_exist() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + let result = db.get_user_variable("room", "username", "myvariable").await; + + assert!(result.is_err()); + assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); + } + + #[tokio::test] + async fn remove_missing_variable_returns_key_does_not_exist() { + let db = Database::new(&tempdir().unwrap()).unwrap(); + let result = db + .delete_user_variable("room", "username", "myvariable") + .await; + + assert!(result.is_err()); + assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); } }