forked from projectmoon/tenebrous-dicebot
Count user variables on a per-room basis.
This commit is contained in:
parent
ebfd230f31
commit
114c879c6f
272
src/db.rs
272
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<LittleEndian>>;
|
||||
|
||||
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<TransactionError<DataError>> for DataError {
|
||||
fn from(error: TransactionError<DataError>) -> 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<u8> {
|
||||
|
@ -42,6 +63,10 @@ fn to_key(room_id: &str, username: &str, variable_name: &str) -> Vec<u8> {
|
|||
key
|
||||
}
|
||||
|
||||
fn room_variable_count_key(room_id: &str, username: &str) -> Vec<u8> {
|
||||
to_key(room_id, username, VARIABLE_COUNT_KEY)
|
||||
}
|
||||
|
||||
fn to_prefix(room_id: &str, username: &str) -> Vec<u8> {
|
||||
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<u8> {
|
|||
prefix
|
||||
}
|
||||
|
||||
fn convert(raw_value: &[u8]) -> Result<i32, DataError> {
|
||||
/// 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<i32, DataError> {
|
||||
let layout = LittleEndianI32Layout::new_unaligned(raw_value.as_ref());
|
||||
|
||||
if let Some(layout) = layout {
|
||||
|
@ -60,6 +87,29 @@ fn convert(raw_value: &[u8]) -> Result<i32, DataError> {
|
|||
}
|
||||
}
|
||||
|
||||
/// 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<i32, DataError> {
|
||||
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<LittleEndian> = I32::new(new_count);
|
||||
variables.insert(key, db_value.as_bytes())?;
|
||||
Ok(new_count)
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub fn new<P: AsRef<std::path::Path>>(path: P) -> Result<Database, DataError> {
|
||||
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<i32, DataError> {
|
||||
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<LittleEndian> = 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<LittleEndian> = 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(_))));
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue