From bf4ce24b79d4cb691e48bf1360d9e6afc22495ac Mon Sep 17 00:00:00 2001 From: projectmoon Date: Tue, 3 Nov 2020 20:14:15 +0000 Subject: [PATCH] Better public API for user variables, avoid allocations in Context. The database API for user variables has changed somewhat again, this time closer to the proper vision. There are now two separate sled Trees in the Variables struct, one for user-defined variables, and one for counts. Keys have been changed to be username-first, then room ID. The signatures of the functions now also use a strongly-typed struct, UserAndRoom. As part of this, the Context object now once again avoids allocating new strings. Other random changes included here: - Remove tempfile crate in favor of sled temporary db config. - Add bincode crate in anticipation of future (de)serializing. --- Cargo.lock | 12 +- Cargo.toml | 6 +- src/bin/dicebot-cmd.rs | 2 +- src/cofd/dice.rs | 24 ++- src/commands.rs | 4 +- src/commands/basic_rolling.rs | 2 +- src/commands/cofd.rs | 2 +- src/commands/cthulhu.rs | 4 +- src/commands/misc.rs | 2 +- src/commands/variables.rs | 33 ++-- src/config.rs | 2 +- src/context.rs | 23 +-- src/db.rs | 19 ++- src/db/data_migrations.rs | 4 +- src/db/errors.rs | 6 + src/db/variables.rs | 279 +++++++++++++++++---------------- src/db/variables/migrations.rs | 157 +++++++++++++++++-- 17 files changed, 368 insertions(+), 213 deletions(-) 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(()) + } +}