diff --git a/Cargo.lock b/Cargo.lock index 60bbe76..0ecf6a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,6 +130,16 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3441f0f7b02788e948e47f457ca01f1d7e6d92c693bc132c22b087d3141c03ff" +[[package]] +name = "bincode" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f30d3a39baa26f9651f17b375061f3233dde33424a8b72b0dbe93a68a0bc896d" +dependencies = [ + "byteorder", + "serde", +] + [[package]] name = "bitflags" version = "1.2.1" @@ -188,6 +198,7 @@ name = "chronicle-dicebot" version = "0.7.0" dependencies = [ "async-trait", + "bincode", "byteorder", "combine", "dirs", @@ -204,7 +215,6 @@ dependencies = [ "rand 0.7.3", "serde", "sled", - "tempfile", "thiserror", "tokio", "toml", diff --git a/Cargo.toml b/Cargo.toml index 32b6968..54a8804 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ zerocopy = "0.3" byteorder = "1.3" futures = "0.3" memmem = "0.1" +bincode = "1.3" phf = { version = "0.7", features = ["macros"] } olm-sys = "1.0" matrix-sdk = { git = "https://github.com/matrix-org/matrix-rust-sdk", rev = "master" } @@ -38,7 +39,4 @@ features = ['derive'] [dependencies.tokio] version = "0.2" -features = [ "full" ] - -[dev-dependencies] -tempfile = "3.1" +features = [ "full" ] \ No newline at end of file diff --git a/src/bin/dicebot-cmd.rs b/src/bin/dicebot-cmd.rs index cdeaf26..a57622a 100644 --- a/src/bin/dicebot-cmd.rs +++ b/src/bin/dicebot-cmd.rs @@ -5,7 +5,7 @@ use chronicle_dicebot::error::BotError; #[tokio::main] async fn main() -> Result<(), BotError> { - let db = Database::new("test-db")?; + let db = Database::new_temp()?; let input = std::env::args().skip(1).collect::>().join(" "); let command = match commands::parse(&input) { Ok(command) => command, diff --git a/src/cofd/dice.rs b/src/cofd/dice.rs index f77d92d..ff645ee 100644 --- a/src/cofd/dice.rs +++ b/src/cofd/dice.rs @@ -1,4 +1,5 @@ use crate::context::Context; +use crate::db::variables::UserAndRoom; use crate::error::BotError; use crate::parser::{Amount, Element, Operator}; use crate::roll::Rolled; @@ -93,15 +94,10 @@ pub struct DicePool { pub(crate) modifiers: DicePoolModifiers, } -async fn calculate_dice_amount<'a>(pool: &'a DicePoolWithContext<'a>) -> Result { +async fn calculate_dice_amount(pool: &DicePoolWithContext<'_>) -> Result { let stream = stream::iter(&pool.0.amounts); - let variables = pool - .1 - .db - .variables - .get_user_variables(&pool.1.room_id, &pool.1.username)?; - - let variables = &variables; + let key = UserAndRoom(&pool.1.username, &pool.1.room_id); + let variables = &pool.1.db.variables.get_user_variables(&key)?; use DiceRollingError::VariableNotFound; let dice_amount: Result = stream @@ -242,7 +238,7 @@ impl DicePoolRoll { } /// Attach a Context to a dice pool. Needed for database access. -pub struct DicePoolWithContext<'a>(pub &'a DicePool, pub &'a Context); +pub struct DicePoolWithContext<'a>(pub &'a DicePool, pub &'a Context<'a>); impl Rolled for DicePoolRoll { fn rolled_value(&self) -> i32 { @@ -368,7 +364,6 @@ pub async fn roll_pool(pool: &DicePoolWithContext<'_>) -> Result Execution; + async fn execute(&self, ctx: &Context<'_>) -> Execution; fn name(&self) -> &'static str; } @@ -64,7 +64,7 @@ pub struct CommandResult { /// go back to Matrix, if the command was executed (successfully or /// not). If a command is determined to be ignored, this function will /// return None, signifying that we should not send a response. -pub async fn execute_command(ctx: &Context) -> Option { +pub async fn execute_command(ctx: &Context<'_>) -> Option { let res = parse(&ctx.message_body); let (plain, html) = match res { diff --git a/src/commands/basic_rolling.rs b/src/commands/basic_rolling.rs index 47d5454..c2e23b5 100644 --- a/src/commands/basic_rolling.rs +++ b/src/commands/basic_rolling.rs @@ -12,7 +12,7 @@ impl Command for RollCommand { "roll regular dice" } - async fn execute(&self, _ctx: &Context) -> Execution { + async fn execute(&self, _ctx: &Context<'_>) -> Execution { let roll = self.0.roll(); let plain = format!("Dice: {}\nResult: {}", self.0, roll); let html = format!( diff --git a/src/commands/cofd.rs b/src/commands/cofd.rs index abed8be..5b74c16 100644 --- a/src/commands/cofd.rs +++ b/src/commands/cofd.rs @@ -11,7 +11,7 @@ impl Command for PoolRollCommand { "roll dice pool" } - async fn execute(&self, ctx: &Context) -> Execution { + async fn execute(&self, ctx: &Context<'_>) -> Execution { let pool_with_ctx = DicePoolWithContext(&self.0, ctx); let roll_result = roll_pool(&pool_with_ctx).await; diff --git a/src/commands/cthulhu.rs b/src/commands/cthulhu.rs index 9f18074..7ccdb24 100644 --- a/src/commands/cthulhu.rs +++ b/src/commands/cthulhu.rs @@ -11,7 +11,7 @@ impl Command for CthRoll { "roll percentile pool" } - async fn execute(&self, _ctx: &Context) -> Execution { + async fn execute(&self, _ctx: &Context<'_>) -> Execution { //TODO this will be converted to a result when supporting variables. let roll = self.0.roll(); let plain = format!("Roll: {}\nResult: {}", self.0, roll); @@ -32,7 +32,7 @@ impl Command for CthAdvanceRoll { "roll percentile pool" } - async fn execute(&self, _ctx: &Context) -> Execution { + async fn execute(&self, _ctx: &Context<'_>) -> Execution { //TODO this will be converted to a result when supporting variables. let roll = self.0.roll(); let plain = format!("Roll: {}\nResult: {}", self.0, roll); diff --git a/src/commands/misc.rs b/src/commands/misc.rs index e78492f..c08230e 100644 --- a/src/commands/misc.rs +++ b/src/commands/misc.rs @@ -11,7 +11,7 @@ impl Command for HelpCommand { "help information" } - async fn execute(&self, _ctx: &Context) -> Execution { + async fn execute(&self, _ctx: &Context<'_>) -> Execution { let help = match &self.0 { Some(topic) => topic.message(), _ => "There is no help for this topic", diff --git a/src/commands/variables.rs b/src/commands/variables.rs index 01a7d10..3f5cb50 100644 --- a/src/commands/variables.rs +++ b/src/commands/variables.rs @@ -1,6 +1,7 @@ use super::{Command, Execution}; use crate::context::Context; use crate::db::errors::DataError; +use crate::db::variables::UserAndRoom; use async_trait::async_trait; pub struct GetAllVariablesCommand; @@ -11,11 +12,9 @@ impl Command for GetAllVariablesCommand { "get all variables" } - async fn execute(&self, ctx: &Context) -> Execution { - let result = ctx - .db - .variables - .get_user_variables(&ctx.room_id, &ctx.username); + async fn execute(&self, ctx: &Context<'_>) -> Execution { + let key = UserAndRoom(&ctx.username, &ctx.room_id); + let result = ctx.db.variables.get_user_variables(&key); let value = match result { Ok(variables) => { @@ -47,12 +46,10 @@ impl Command for GetVariableCommand { "retrieve variable value" } - async fn execute(&self, ctx: &Context) -> Execution { + async fn execute(&self, ctx: &Context<'_>) -> Execution { let name = &self.0; - let result = ctx - .db - .variables - .get_user_variable(&ctx.room_id, &ctx.username, name); + let key = UserAndRoom(&ctx.username, &ctx.room_id); + let result = ctx.db.variables.get_user_variable(&key, name); let value = match result { Ok(num) => format!("{} = {}", name, num), @@ -74,13 +71,11 @@ impl Command for SetVariableCommand { "set variable value" } - async fn execute(&self, ctx: &Context) -> Execution { + async fn execute(&self, ctx: &Context<'_>) -> Execution { let name = &self.0; let value = self.1; - let result = ctx - .db - .variables - .set_user_variable(&ctx.room_id, &ctx.username, name, value); + let key = UserAndRoom(&ctx.username, ctx.room_id); + let result = ctx.db.variables.set_user_variable(&key, name, value); let content = match result { Ok(_) => format!("{} = {}", name, value), @@ -101,12 +96,10 @@ impl Command for DeleteVariableCommand { "delete variable" } - async fn execute(&self, ctx: &Context) -> Execution { + async fn execute(&self, ctx: &Context<'_>) -> Execution { let name = &self.0; - let result = ctx - .db - .variables - .delete_user_variable(&ctx.room_id, &ctx.username, name); + let key = UserAndRoom(&ctx.username, ctx.room_id); + let result = ctx.db.variables.delete_user_variable(&key, name); let value = match result { Ok(()) => format!("{} now unset", name), diff --git a/src/config.rs b/src/config.rs index 0d7dc2a..22e9b91 100644 --- a/src/config.rs +++ b/src/config.rs @@ -6,7 +6,7 @@ use thiserror::Error; /// Shortcut to defining db migration versions. Will probably /// eventually be moved to a config file. -const MIGRATION_VERSION: u32 = 4; +const MIGRATION_VERSION: u32 = 5; #[derive(Error, Debug)] pub enum ConfigError { diff --git a/src/context.rs b/src/context.rs index d40623c..630659f 100644 --- a/src/context.rs +++ b/src/context.rs @@ -3,20 +3,25 @@ use crate::db::Database; /// A context carried through the system providing access to things /// like the database. #[derive(Clone)] -pub struct Context { +pub struct Context<'a> { pub db: Database, - pub room_id: String, - pub username: String, - pub message_body: String, + pub room_id: &'a str, + pub username: &'a str, + pub message_body: &'a str, } -impl Context { - pub fn new(db: &Database, room_id: &str, username: &str, message_body: &str) -> Context { +impl<'a> Context<'a> { + pub fn new( + db: &Database, + room_id: &'a str, + username: &'a str, + message_body: &'a str, + ) -> Context<'a> { Context { db: db.clone(), - room_id: room_id.to_owned(), - username: username.to_owned(), - message_body: message_body.to_owned(), + room_id: room_id, + username: username, + message_body: message_body, } } } diff --git a/src/db.rs b/src/db.rs index 532153f..62a94e6 100644 --- a/src/db.rs +++ b/src/db.rs @@ -2,7 +2,7 @@ use crate::db::errors::{DataError, MigrationError}; use crate::db::migrations::{get_migration_version, Migrations}; use crate::db::variables::Variables; use log::info; -use sled::Db; +use sled::{Config, Db}; use std::path::Path; pub mod data_migrations; @@ -19,18 +19,27 @@ pub struct Database { } impl Database { - pub fn new>(path: P) -> Result { - let db = sled::open(path)?; - let variables = db.open_tree("variables")?; + fn new_db(db: sled::Db) -> Result { let migrations = db.open_tree("migrations")?; Ok(Database { db: db.clone(), - variables: Variables(variables), + variables: Variables::new(&db)?, migrations: Migrations(migrations), }) } + pub fn new>(path: P) -> Result { + let db = sled::open(path)?; + Self::new_db(db) + } + + pub fn new_temp() -> Result { + let config = Config::new().temporary(true); + let db = config.open()?; + Self::new_db(db) + } + pub fn migrate(&self, to_version: u32) -> Result<(), DataError> { //get version from db let db_version = get_migration_version(&self)?; diff --git a/src/db/data_migrations.rs b/src/db/data_migrations.rs index 3ab674a..3669ddb 100644 --- a/src/db/data_migrations.rs +++ b/src/db/data_migrations.rs @@ -9,7 +9,9 @@ static MIGRATIONS: phf::Map = phf_map! { 1u32 => (add_room_user_variable_count::migrate, "add_room_user_variable_count"), 2u32 => (delete_v0_schema, "delete_v0_schema"), 3u32 => (delete_variable_count, "delete_variable_count"), - 4u32 => (change_delineator_delimiter::migrate, "change_delineator_delimiter") + 4u32 => (change_delineator_delimiter::migrate, "change_delineator_delimiter"), + 5u32 => (change_tree_structure::migrate, "change_tree_structure"), + }; pub fn get_migrations(versions: &[u32]) -> Result, MigrationError> { diff --git a/src/db/errors.rs b/src/db/errors.rs index b1f9ab0..cb25c9b 100644 --- a/src/db/errors.rs +++ b/src/db/errors.rs @@ -20,6 +20,9 @@ pub enum DataError { #[error("value does not exist for key: {0}")] KeyDoesNotExist(String), + #[error("too many entries")] + TooManyEntries, + #[error("expected i32, but i32 schema was violated")] I32SchemaViolation, @@ -37,6 +40,9 @@ pub enum DataError { #[error("data migration error: {0}")] MigrationError(#[from] MigrationError), + + #[error("deserialization error: {0}")] + DeserializationError(#[from] bincode::Error), } /// This From implementation is necessary to deal with the recursive diff --git a/src/db/variables.rs b/src/db/variables.rs index ae2d016..8d2be34 100644 --- a/src/db/variables.rs +++ b/src/db/variables.rs @@ -2,6 +2,7 @@ use crate::db::errors::DataError; use crate::db::schema::convert_i32; use byteorder::LittleEndian; use sled::transaction::{abort, TransactionalTree}; +use sled::Transactional; use sled::Tree; use std::collections::HashMap; use std::convert::From; @@ -11,74 +12,48 @@ use zerocopy::AsBytes; pub(super) mod migrations; -const METADATA_SPACE: &'static [u8] = b"metadata"; -const VARIABLE_SPACE: &'static [u8] = b"variables"; - -const VARIABLE_COUNT_KEY: &'static str = "variable_count"; - #[derive(Clone)] -pub struct Variables(pub(super) Tree); +pub struct Variables { + //room id + username + variable = i32 + pub(in crate::db) room_user_variables: Tree, -//TODO at least some of these will probalby move elsewhere. + //room id + username = i32 + pub(in crate::db) room_user_variable_count: Tree, +} -fn space_prefix>>(space: &[u8], delineator: D) -> Vec { - let mut metadata_prefix = vec![]; - metadata_prefix.extend_from_slice(space); - metadata_prefix.push(0xff); - let delineator = delineator.into(); +/// Request soemthing by a username and room ID. +pub struct UserAndRoom<'a>(pub &'a str, pub &'a str); - if delineator.len() > 0 { - metadata_prefix.extend_from_slice(delineator.as_bytes()); - metadata_prefix.push(0xff); +fn to_vec(value: &UserAndRoom<'_>) -> Vec { + let mut bytes = vec![]; + bytes.extend_from_slice(value.0.as_bytes()); + bytes.push(0xfe); + bytes.extend_from_slice(value.1.as_bytes()); + bytes +} + +impl From> for Vec { + fn from(value: UserAndRoom) -> Vec { + to_vec(&value) } - - metadata_prefix } -fn metadata_space_prefix>>(delineator: D) -> Vec { - space_prefix(METADATA_SPACE, delineator) -} - -fn metadata_space_key>>(delineator: D, key_name: &str) -> Vec { - let mut metadata_key = metadata_space_prefix(delineator); - metadata_key.extend_from_slice(key_name.as_bytes()); - metadata_key -} - -fn variables_space_prefix>>(delineator: D) -> Vec { - space_prefix(VARIABLE_SPACE, delineator) -} - -fn variables_space_key>>(delineator: D, key_name: &str) -> Vec { - let mut metadata_key = variables_space_prefix(delineator); - metadata_key.extend_from_slice(key_name.as_bytes()); - metadata_key -} - -/// Delineator for keeping track of a key by room ID and username. -struct RoomAndUser<'a>(&'a str, &'a str); - -impl<'a> From> for Vec { - fn from(value: RoomAndUser<'a>) -> Vec { - let mut bytes = vec![]; - bytes.extend_from_slice(value.0.as_bytes()); - bytes.push(0xfe); - bytes.extend_from_slice(value.1.as_bytes()); - bytes +impl From<&UserAndRoom<'_>> for Vec { + fn from(value: &UserAndRoom) -> Vec { + to_vec(value) } } /// Use a transaction to 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, + room_variable_count: &TransactionalTree, + user_and_room: &UserAndRoom<'_>, amount: i32, ) -> Result { - let key = metadata_space_key(RoomAndUser(room_id, username), VARIABLE_COUNT_KEY); + let key: Vec = user_and_room.into(); - let mut new_count = match variables.get(&key)? { + let mut new_count = match room_variable_count.get(&key)? { Some(bytes) => convert_i32(&bytes)? + amount, None => amount, }; @@ -88,21 +63,28 @@ fn alter_room_variable_count( } let db_value: I32 = I32::new(new_count); - variables.insert(key, db_value.as_bytes())?; + room_variable_count.insert(key, db_value.as_bytes())?; Ok(new_count) } impl Variables { + pub(in crate::db) fn new(db: &sled::Db) -> Result { + Ok(Variables { + room_user_variables: db.open_tree("variables")?, + room_user_variable_count: db.open_tree("room_user_variable_count")?, + }) + } + pub fn get_user_variables( &self, - room_id: &str, - username: &str, + key: &UserAndRoom<'_>, ) -> Result, DataError> { - let prefix = variables_space_prefix(RoomAndUser(room_id, username)); + let mut prefix: Vec = key.into(); + prefix.push(0xff); let prefix_len: usize = prefix.len(); let variables: Result, DataError> = self - .0 + .room_user_variables .scan_prefix(prefix) .map(|entry| match entry { Ok((key, raw_value)) => { @@ -114,99 +96,104 @@ impl Variables { }) .collect(); - //Convert I32 to hash map. Can we do this in the first mapping - //step instead? For some reason this is faster. + //Convert I32 to hash map. collect() inferred via return type. variables.map(|entries| entries.into_iter().collect()) } - pub fn get_variable_count(&self, room_id: &str, username: &str) -> Result { - let delineator = RoomAndUser(room_id, username); - let key = metadata_space_key(delineator, VARIABLE_COUNT_KEY); + pub fn get_variable_count(&self, user_and_room: &UserAndRoom<'_>) -> Result { + let key: Vec = user_and_room.into(); - if let Some(raw_value) = self.0.get(&key)? { - convert_i32(&raw_value) - } else { - Ok(0) + match self.room_user_variable_count.get(&key)? { + Some(raw_value) => convert_i32(&raw_value), + None => Ok(0), } } pub fn get_user_variable( &self, - room_id: &str, - username: &str, + user_and_room: &UserAndRoom<'_>, variable_name: &str, ) -> Result { - let key = variables_space_key(RoomAndUser(room_id, username), variable_name); + let mut key: Vec = user_and_room.into(); + key.push(0xff); + key.extend_from_slice(variable_name.as_bytes()); - if let Some(raw_value) = self.0.get(&key)? { - convert_i32(&raw_value) - } else { - Err(DataError::KeyDoesNotExist(variable_name.to_owned())) + match self.room_user_variables.get(&key)? { + Some(raw_value) => convert_i32(&raw_value), + _ => Err(DataError::KeyDoesNotExist(variable_name.to_owned())), } } pub fn set_user_variable( &self, - room_id: &str, - username: &str, + user_and_room: &UserAndRoom<'_>, variable_name: &str, value: i32, ) -> Result<(), DataError> { - self.0 - .transaction(|tx| { - let key = variables_space_key(RoomAndUser(room_id, username), variable_name); + if self.get_variable_count(user_and_room)? >= 100 { + return Err(DataError::TooManyEntries); + } + + (&self.room_user_variables, &self.room_user_variable_count).transaction( + |(tx_vars, tx_counts)| { + let mut key: Vec = user_and_room.into(); + key.push(0xff); + key.extend_from_slice(variable_name.as_bytes()); + let db_value: I32 = I32::new(value); - let old_value = tx.insert(key, db_value.as_bytes())?; + let old_value = tx_vars.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(()), + if let Err(e) = alter_room_variable_count(&tx_counts, &user_and_room, 1) { + return abort(e); } - } else { - Ok(()) } - }) - .map_err(|e| e.into()) + + Ok(()) + }, + )?; + + Ok(()) } pub fn delete_user_variable( &self, - room_id: &str, - username: &str, + user_and_room: &UserAndRoom<'_>, variable_name: &str, ) -> Result<(), DataError> { - self.0 - .transaction(|tx| { - let key = variables_space_key(RoomAndUser(room_id, username), variable_name); + (&self.room_user_variables, &self.room_user_variable_count).transaction( + |(tx_vars, tx_counts)| { + let mut key: Vec = user_and_room.into(); + key.push(0xff); + key.extend_from_slice(variable_name.as_bytes()); //TODO why does tx.remove require moving the key? - if let Some(_) = tx.remove(key.clone())? { - match alter_room_variable_count(&tx, room_id, username, -1) { - Err(e) => abort(e), - _ => Ok(()), + if let Some(_) = tx_vars.remove(key.clone())? { + if let Err(e) = alter_room_variable_count(&tx_counts, user_and_room, -1) { + return abort(e); } } else { - abort(DataError::KeyDoesNotExist(variable_name.to_owned())) + return abort(DataError::KeyDoesNotExist(variable_name.to_owned())); } - }) - .map_err(|e| e.into()) + + Ok(()) + }, + )?; + + Ok(()) } } #[cfg(test)] mod tests { use super::*; - use tempfile::tempdir; + use sled::Config; fn create_test_instance() -> Variables { - Variables( - sled::open(&tempdir().unwrap()) - .unwrap() - .open_tree("variables") - .unwrap(), - ) + let config = Config::new().temporary(true); + let db = config.open().unwrap(); + Variables::new(&db).unwrap() } //Room Variable count tests @@ -214,24 +201,23 @@ mod tests { #[test] fn alter_room_variable_count_test() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); let alter_count = |amount: i32| { variables - .0 - .transaction(|tx| { - match alter_room_variable_count(&tx, "room", "username", amount) { - Err(e) => abort(e), - _ => Ok(()), - } + .room_user_variable_count + .transaction(|tx| match alter_room_variable_count(&tx, &key, amount) { + Err(e) => abort(e), + _ => Ok(()), }) .expect("got transaction failure"); }; - fn get_count(variables: &Variables) -> i32 { + let get_count = |variables: &Variables| -> i32 { variables - .get_variable_count("room", "username") + .get_variable_count(&key) .expect("could not get variable count") - } + }; //addition alter_count(5); @@ -245,19 +231,18 @@ mod tests { #[test] fn alter_room_variable_count_cannot_go_below_0_test() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); variables - .0 - .transaction( - |tx| match alter_room_variable_count(&tx, "room", "username", -1000) { - Err(e) => abort(e), - _ => Ok(()), - }, - ) + .room_user_variable_count + .transaction(|tx| match alter_room_variable_count(&tx, &key, -1000) { + Err(e) => abort(e), + _ => Ok(()), + }) .expect("got transaction failure"); let count = variables - .get_variable_count("room", "username") + .get_variable_count(&key) .expect("could not get variable count"); assert_eq!(0, count); @@ -266,9 +251,10 @@ mod tests { #[test] fn empty_db_reports_0_room_variable_count_test() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); let count = variables - .get_variable_count("room", "username") + .get_variable_count(&key) .expect("could not get variable count"); assert_eq!(0, count); @@ -277,13 +263,14 @@ mod tests { #[test] fn set_user_variable_increments_count() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); variables - .set_user_variable("room", "username", "myvariable", 5) + .set_user_variable(&key, "myvariable", 5) .expect("could not insert variable"); let count = variables - .get_variable_count("room", "username") + .get_variable_count(&key) .expect("could not get variable count"); assert_eq!(1, count); @@ -292,17 +279,18 @@ mod tests { #[test] fn update_user_variable_does_not_increment_count() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); variables - .set_user_variable("room", "username", "myvariable", 5) + .set_user_variable(&key, "myvariable", 5) .expect("could not insert variable"); variables - .set_user_variable("room", "username", "myvariable", 10) + .set_user_variable(&key, "myvariable", 10) .expect("could not update variable"); let count = variables - .get_variable_count("room", "username") + .get_variable_count(&key) .expect("could not get variable count"); assert_eq!(1, count); @@ -313,30 +301,49 @@ mod tests { #[test] fn set_and_get_variable_test() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); + variables - .set_user_variable("room", "username", "myvariable", 5) + .set_user_variable(&key, "myvariable", 5) .expect("could not insert variable"); let value = variables - .get_user_variable("room", "username", "myvariable") + .get_user_variable(&key, "myvariable") .expect("could not get value"); assert_eq!(5, value); } + #[test] + fn cannot_set_more_than_100_variables_per_room() { + let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); + + for c in 0..100 { + variables + .set_user_variable(&key, &format!("myvariable{}", c), 5) + .expect("could not insert variable"); + } + + let result = variables.set_user_variable(&key, "myvariable101", 5); + assert!(result.is_err()); + assert!(matches!(result, Err(DataError::TooManyEntries))); + } + #[test] fn delete_variable_test() { let variables = create_test_instance(); + let key = UserAndRoom("username", "room"); variables - .set_user_variable("room", "username", "myvariable", 5) + .set_user_variable(&key, "myvariable", 5) .expect("could not insert variable"); variables - .delete_user_variable("room", "username", "myvariable") + .delete_user_variable(&key, "myvariable") .expect("could not delete value"); - let result = variables.get_user_variable("room", "username", "myvariable"); + let result = variables.get_user_variable(&key, "myvariable"); assert!(result.is_err()); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); @@ -345,7 +352,8 @@ mod tests { #[test] fn get_missing_variable_returns_key_does_not_exist() { let variables = create_test_instance(); - let result = variables.get_user_variable("room", "username", "myvariable"); + let key = UserAndRoom("username", "room"); + let result = variables.get_user_variable(&key, "myvariable"); assert!(result.is_err()); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); @@ -354,7 +362,8 @@ mod tests { #[test] fn remove_missing_variable_returns_key_does_not_exist() { let variables = create_test_instance(); - let result = variables.delete_user_variable("room", "username", "myvariable"); + let key = UserAndRoom("username", "room"); + let result = variables.delete_user_variable(&key, "myvariable"); assert!(result.is_err()); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); diff --git a/src/db/variables/migrations.rs b/src/db/variables/migrations.rs index a18a1c2..6db15e9 100644 --- a/src/db/variables/migrations.rs +++ b/src/db/variables/migrations.rs @@ -5,7 +5,8 @@ use byteorder::LittleEndian; use memmem::{Searcher, TwoWaySearcher}; use sled::transaction::TransactionError; use sled::{Batch, IVec}; -use zerocopy::byteorder::U32; +use std::collections::HashMap; +use zerocopy::byteorder::{I32, U32}; use zerocopy::AsBytes; pub(in crate::db) mod add_room_user_variable_count { @@ -52,9 +53,20 @@ pub(in crate::db) mod add_room_user_variable_count { } } + fn create_key(room_id: &str, username: &str) -> Vec { + let mut key = b"variables".to_vec(); + key.push(0xff); + key.extend_from_slice(room_id.as_bytes()); + key.push(0xff); + key.extend_from_slice(username.as_bytes()); + key.push(0xff); + key.extend_from_slice(b"variable_count"); + key + } + pub(in crate::db) fn migrate(db: &Database) -> Result<(), DataError> { - let tree = &db.variables.0; - let prefix = variables_space_prefix(""); + let tree = &db.variables.room_user_variables; + let prefix = b"variables"; //Extract a vec of tuples, consisting of room id + username. let results: Vec = tree @@ -73,13 +85,12 @@ pub(in crate::db) mod add_room_user_variable_count { //Start a transaction on the variables tree. let tx_result: Result<_, TransactionError> = - db.variables.0.transaction(|tx_vars| { + db.variables.room_user_variables.transaction(|tx_vars| { let batch = counts.iter().fold(Batch::default(), |mut batch, entry| { let (info, count) = entry; //Add variable count according to new schema. - let delineator = super::RoomAndUser(&info.room_id, &info.username); - let key = variables_space_key(delineator, VARIABLE_COUNT_KEY); + let key = create_key(&info.room_id, &info.username); let db_value: U32 = U32::new(*count); batch.insert(key, db_value.as_bytes()); @@ -99,7 +110,7 @@ pub(in crate::db) mod add_room_user_variable_count { } pub(in crate::db) fn delete_v0_schema(db: &Database) -> Result<(), DataError> { - let mut vars = db.variables.0.scan_prefix(""); + let mut vars = db.variables.room_user_variables.scan_prefix(""); let mut batch = Batch::default(); while let Some(Ok((key, _))) = vars.next() { @@ -110,13 +121,13 @@ pub(in crate::db) fn delete_v0_schema(db: &Database) -> Result<(), DataError> { } } - db.variables.0.apply_batch(batch)?; + db.variables.room_user_variables.apply_batch(batch)?; Ok(()) } pub(in crate::db) fn delete_variable_count(db: &Database) -> Result<(), DataError> { - let prefix = variables_space_prefix(""); - let mut vars = db.variables.0.scan_prefix(prefix); + let prefix = b"variables"; + let mut vars = db.variables.room_user_variables.scan_prefix(prefix); let mut batch = Batch::default(); while let Some(Ok((key, _))) = vars.next() { @@ -133,7 +144,7 @@ pub(in crate::db) fn delete_variable_count(db: &Database) -> Result<(), DataErro } } - db.variables.0.apply_batch(batch)?; + db.variables.room_user_variables.apply_batch(batch)?; Ok(()) } @@ -203,8 +214,8 @@ pub(in crate::db) mod change_delineator_delimiter { } pub fn migrate(db: &Database) -> Result<(), DataError> { - let tree = &db.variables.0; - let prefix = variables_space_prefix(""); + let tree = &db.variables.room_user_variables; + let prefix = b"variables"; let results: Vec = tree .scan_prefix(&prefix) @@ -214,8 +225,8 @@ pub(in crate::db) mod change_delineator_delimiter { let mut batch = Batch::default(); for insert in results { - let old = create_old_key(&prefix, &insert); - let new = create_new_key(&prefix, &insert); + let old = create_old_key(prefix, &insert); + let new = create_new_key(prefix, &insert); batch.remove(old); batch.insert(new, insert.value); @@ -225,3 +236,119 @@ pub(in crate::db) mod change_delineator_delimiter { Ok(()) } } + +/// Move the user variable entries into two tree structures, with yet +/// another key format change. Now there is one tree for variable +/// counts, and one tree for actual user variables. Keys in the user +/// variable tree were changed to be username-first, then room ID. +/// They are still separated by 0xfe, while the variable name is +/// separated by 0xff. Variable count now stores just +/// USERNAME0xfeROOM_ID and a count in its own tree. This enables +/// public use of a strongly typed UserAndRoom struct for getting +/// variables. +pub(in crate::db) mod change_tree_structure { + use super::*; + + /// An entry in the room user variables keyspace. + struct UserVariableEntry { + room_id: String, + username: String, + variable_name: String, + value: IVec, + } + + /// Extract keys and values from the variables keyspace according + /// to the v1 schema. + fn extract_v1_entries( + entry: sled::Result<(IVec, IVec)>, + ) -> Result { + if let Ok((key, value)) = entry { + let keys: Vec> = key + .split(|&b| b == 0xff || b == 0xfe) + .map(|b| str::from_utf8(b)) + .collect(); + + if let &[_, Ok(room_id), Ok(username), Ok(variable)] = keys.as_slice() { + Ok(UserVariableEntry { + room_id: room_id.to_owned(), + username: username.to_owned(), + variable_name: variable.to_owned(), + value: value, + }) + } else { + Err(MigrationError::MigrationFailed( + "a key violates utf8 schema".to_string(), + )) + } + } else { + Err(MigrationError::MigrationFailed( + "encountered unexpected key".to_string(), + )) + } + } + + /// Create an old key, of "variables" 0xff "room id" 0xfe "username" 0xff "variablename". + fn create_old_key(prefix: &[u8], insert: &UserVariableEntry) -> Vec { + let mut key = vec![]; + key.extend_from_slice(&prefix); //prefix already has 0xff. + key.extend_from_slice(&insert.room_id.as_bytes()); + key.push(0xff); + key.extend_from_slice(&insert.username.as_bytes()); + key.push(0xff); + key.extend_from_slice(&insert.variable_name.as_bytes()); + key + } + + /// Create a new key, of "username" 0xfe "room id" 0xff "variablename". + fn create_new_key(insert: &UserVariableEntry) -> Vec { + let mut key = vec![]; + key.extend_from_slice(&insert.username.as_bytes()); + key.push(0xfe); + key.extend_from_slice(&insert.room_id.as_bytes()); + key.push(0xff); + key.extend_from_slice(&insert.variable_name.as_bytes()); + key + } + + pub fn migrate(db: &Database) -> Result<(), DataError> { + let variables_tree = &db.variables.room_user_variables; + let count_tree = &db.variables.room_user_variable_count; + let prefix = b"variables"; + + let results: Vec = variables_tree + .scan_prefix(&prefix) + .map(extract_v1_entries) + .collect::, MigrationError>>()?; + + let mut counts: HashMap<(String, String), i32> = HashMap::new(); + let mut batch = Batch::default(); + + for insert in results { + let count = counts + .entry((insert.username.clone(), insert.room_id.clone())) + .or_insert(0); + *count += 1; + + let old = create_old_key(prefix, &insert); + let new = create_new_key(&insert); + + batch.remove(old); + batch.insert(new, insert.value); + } + + let mut count_batch = Batch::default(); + counts.into_iter().for_each(|((username, room_id), count)| { + let mut key = username.as_bytes().to_vec(); + key.push(0xfe); + key.extend_from_slice(room_id.as_bytes()); + + let db_value: I32 = I32::new(count); + count_batch.insert(key, db_value.as_bytes()); + }); + + variables_tree.apply_batch(batch)?; + count_tree.apply_batch(count_batch)?; + + Ok(()) + } +}