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.
This commit is contained in:
projectmoon 2020-11-03 20:14:15 +00:00 committed by ProjectMoon
parent 8ec19c266f
commit bf4ce24b79
17 changed files with 368 additions and 213 deletions

12
Cargo.lock generated
View File

@ -130,6 +130,16 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3441f0f7b02788e948e47f457ca01f1d7e6d92c693bc132c22b087d3141c03ff" 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]] [[package]]
name = "bitflags" name = "bitflags"
version = "1.2.1" version = "1.2.1"
@ -188,6 +198,7 @@ name = "chronicle-dicebot"
version = "0.7.0" version = "0.7.0"
dependencies = [ dependencies = [
"async-trait", "async-trait",
"bincode",
"byteorder", "byteorder",
"combine", "combine",
"dirs", "dirs",
@ -204,7 +215,6 @@ dependencies = [
"rand 0.7.3", "rand 0.7.3",
"serde", "serde",
"sled", "sled",
"tempfile",
"thiserror", "thiserror",
"tokio", "tokio",
"toml", "toml",

View File

@ -28,6 +28,7 @@ zerocopy = "0.3"
byteorder = "1.3" byteorder = "1.3"
futures = "0.3" futures = "0.3"
memmem = "0.1" memmem = "0.1"
bincode = "1.3"
phf = { version = "0.7", features = ["macros"] } phf = { version = "0.7", features = ["macros"] }
olm-sys = "1.0" olm-sys = "1.0"
matrix-sdk = { git = "https://github.com/matrix-org/matrix-rust-sdk", rev = "master" } matrix-sdk = { git = "https://github.com/matrix-org/matrix-rust-sdk", rev = "master" }
@ -38,7 +39,4 @@ features = ['derive']
[dependencies.tokio] [dependencies.tokio]
version = "0.2" version = "0.2"
features = [ "full" ] features = [ "full" ]
[dev-dependencies]
tempfile = "3.1"

View File

@ -5,7 +5,7 @@ use chronicle_dicebot::error::BotError;
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), BotError> { async fn main() -> Result<(), BotError> {
let db = Database::new("test-db")?; let db = Database::new_temp()?;
let input = std::env::args().skip(1).collect::<Vec<String>>().join(" "); let input = std::env::args().skip(1).collect::<Vec<String>>().join(" ");
let command = match commands::parse(&input) { let command = match commands::parse(&input) {
Ok(command) => command, Ok(command) => command,

View File

@ -1,4 +1,5 @@
use crate::context::Context; use crate::context::Context;
use crate::db::variables::UserAndRoom;
use crate::error::BotError; use crate::error::BotError;
use crate::parser::{Amount, Element, Operator}; use crate::parser::{Amount, Element, Operator};
use crate::roll::Rolled; use crate::roll::Rolled;
@ -93,15 +94,10 @@ pub struct DicePool {
pub(crate) modifiers: DicePoolModifiers, pub(crate) modifiers: DicePoolModifiers,
} }
async fn calculate_dice_amount<'a>(pool: &'a DicePoolWithContext<'a>) -> Result<i32, BotError> { async fn calculate_dice_amount(pool: &DicePoolWithContext<'_>) -> Result<i32, BotError> {
let stream = stream::iter(&pool.0.amounts); let stream = stream::iter(&pool.0.amounts);
let variables = pool let key = UserAndRoom(&pool.1.username, &pool.1.room_id);
.1 let variables = &pool.1.db.variables.get_user_variables(&key)?;
.db
.variables
.get_user_variables(&pool.1.room_id, &pool.1.username)?;
let variables = &variables;
use DiceRollingError::VariableNotFound; use DiceRollingError::VariableNotFound;
let dice_amount: Result<i32, BotError> = stream let dice_amount: Result<i32, BotError> = stream
@ -242,7 +238,7 @@ impl DicePoolRoll {
} }
/// Attach a Context to a dice pool. Needed for database access. /// 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 { impl Rolled for DicePoolRoll {
fn rolled_value(&self) -> i32 { fn rolled_value(&self) -> i32 {
@ -368,7 +364,6 @@ pub async fn roll_pool(pool: &DicePoolWithContext<'_>) -> Result<RolledDicePool,
mod tests { mod tests {
use super::*; use super::*;
use crate::db::Database; use crate::db::Database;
use tempfile::tempdir;
///Instead of being random, generate a series of numbers we have complete ///Instead of being random, generate a series of numbers we have complete
///control over. ///control over.
@ -507,7 +502,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn rejects_large_expression_test() { async fn rejects_large_expression_test() {
let db = Database::new(&tempdir().unwrap()).unwrap(); let db = Database::new_temp().unwrap();
let ctx = Context::new(&db, "roomid", "username", "message"); let ctx = Context::new(&db, "roomid", "username", "message");
let mut amounts = vec![]; let mut amounts = vec![];
@ -532,7 +527,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn converts_to_chance_die_test() { async fn converts_to_chance_die_test() {
let db = Database::new(&tempdir().unwrap()).unwrap(); let db = Database::new_temp().unwrap();
let ctx = Context::new(&db, "roomid", "username", "message"); let ctx = Context::new(&db, "roomid", "username", "message");
let mut amounts = vec![]; let mut amounts = vec![];
@ -554,11 +549,12 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn can_resolve_variables_test() { async fn can_resolve_variables_test() {
let db = Database::new(&tempdir().unwrap()).unwrap(); let db = Database::new_temp().unwrap();
let ctx = Context::new(&db, "roomid", "username", "message"); let ctx = Context::new(&db, "roomid", "username", "message");
let user_and_room = crate::db::variables::UserAndRoom(&ctx.username, &ctx.room_id);
db.variables db.variables
.set_user_variable(&ctx.room_id, &ctx.username, "myvariable", 10) .set_user_variable(&user_and_room, "myvariable", 10)
.expect("could not set myvariable to 10"); .expect("could not set myvariable to 10");
let amounts = vec![Amount { let amounts = vec![Amount {

View File

@ -37,7 +37,7 @@ impl Execution {
#[async_trait] #[async_trait]
pub trait Command: Send + Sync { pub trait Command: Send + Sync {
async fn execute(&self, ctx: &Context) -> Execution; async fn execute(&self, ctx: &Context<'_>) -> Execution;
fn name(&self) -> &'static str; fn name(&self) -> &'static str;
} }
@ -64,7 +64,7 @@ pub struct CommandResult {
/// go back to Matrix, if the command was executed (successfully or /// go back to Matrix, if the command was executed (successfully or
/// not). If a command is determined to be ignored, this function will /// not). If a command is determined to be ignored, this function will
/// return None, signifying that we should not send a response. /// return None, signifying that we should not send a response.
pub async fn execute_command(ctx: &Context) -> Option<CommandResult> { pub async fn execute_command(ctx: &Context<'_>) -> Option<CommandResult> {
let res = parse(&ctx.message_body); let res = parse(&ctx.message_body);
let (plain, html) = match res { let (plain, html) = match res {

View File

@ -12,7 +12,7 @@ impl Command for RollCommand {
"roll regular dice" "roll regular dice"
} }
async fn execute(&self, _ctx: &Context) -> Execution { async fn execute(&self, _ctx: &Context<'_>) -> Execution {
let roll = self.0.roll(); let roll = self.0.roll();
let plain = format!("Dice: {}\nResult: {}", self.0, roll); let plain = format!("Dice: {}\nResult: {}", self.0, roll);
let html = format!( let html = format!(

View File

@ -11,7 +11,7 @@ impl Command for PoolRollCommand {
"roll dice pool" "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 pool_with_ctx = DicePoolWithContext(&self.0, ctx);
let roll_result = roll_pool(&pool_with_ctx).await; let roll_result = roll_pool(&pool_with_ctx).await;

View File

@ -11,7 +11,7 @@ impl Command for CthRoll {
"roll percentile pool" "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. //TODO this will be converted to a result when supporting variables.
let roll = self.0.roll(); let roll = self.0.roll();
let plain = format!("Roll: {}\nResult: {}", self.0, roll); let plain = format!("Roll: {}\nResult: {}", self.0, roll);
@ -32,7 +32,7 @@ impl Command for CthAdvanceRoll {
"roll percentile pool" "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. //TODO this will be converted to a result when supporting variables.
let roll = self.0.roll(); let roll = self.0.roll();
let plain = format!("Roll: {}\nResult: {}", self.0, roll); let plain = format!("Roll: {}\nResult: {}", self.0, roll);

View File

@ -11,7 +11,7 @@ impl Command for HelpCommand {
"help information" "help information"
} }
async fn execute(&self, _ctx: &Context) -> Execution { async fn execute(&self, _ctx: &Context<'_>) -> Execution {
let help = match &self.0 { let help = match &self.0 {
Some(topic) => topic.message(), Some(topic) => topic.message(),
_ => "There is no help for this topic", _ => "There is no help for this topic",

View File

@ -1,6 +1,7 @@
use super::{Command, Execution}; use super::{Command, Execution};
use crate::context::Context; use crate::context::Context;
use crate::db::errors::DataError; use crate::db::errors::DataError;
use crate::db::variables::UserAndRoom;
use async_trait::async_trait; use async_trait::async_trait;
pub struct GetAllVariablesCommand; pub struct GetAllVariablesCommand;
@ -11,11 +12,9 @@ impl Command for GetAllVariablesCommand {
"get all variables" "get all variables"
} }
async fn execute(&self, ctx: &Context) -> Execution { async fn execute(&self, ctx: &Context<'_>) -> Execution {
let result = ctx let key = UserAndRoom(&ctx.username, &ctx.room_id);
.db let result = ctx.db.variables.get_user_variables(&key);
.variables
.get_user_variables(&ctx.room_id, &ctx.username);
let value = match result { let value = match result {
Ok(variables) => { Ok(variables) => {
@ -47,12 +46,10 @@ impl Command for GetVariableCommand {
"retrieve variable value" "retrieve variable value"
} }
async fn execute(&self, ctx: &Context) -> Execution { async fn execute(&self, ctx: &Context<'_>) -> Execution {
let name = &self.0; let name = &self.0;
let result = ctx let key = UserAndRoom(&ctx.username, &ctx.room_id);
.db let result = ctx.db.variables.get_user_variable(&key, name);
.variables
.get_user_variable(&ctx.room_id, &ctx.username, name);
let value = match result { let value = match result {
Ok(num) => format!("{} = {}", name, num), Ok(num) => format!("{} = {}", name, num),
@ -74,13 +71,11 @@ impl Command for SetVariableCommand {
"set variable value" "set variable value"
} }
async fn execute(&self, ctx: &Context) -> Execution { async fn execute(&self, ctx: &Context<'_>) -> Execution {
let name = &self.0; let name = &self.0;
let value = self.1; let value = self.1;
let result = ctx let key = UserAndRoom(&ctx.username, ctx.room_id);
.db let result = ctx.db.variables.set_user_variable(&key, name, value);
.variables
.set_user_variable(&ctx.room_id, &ctx.username, name, value);
let content = match result { let content = match result {
Ok(_) => format!("{} = {}", name, value), Ok(_) => format!("{} = {}", name, value),
@ -101,12 +96,10 @@ impl Command for DeleteVariableCommand {
"delete variable" "delete variable"
} }
async fn execute(&self, ctx: &Context) -> Execution { async fn execute(&self, ctx: &Context<'_>) -> Execution {
let name = &self.0; let name = &self.0;
let result = ctx let key = UserAndRoom(&ctx.username, ctx.room_id);
.db let result = ctx.db.variables.delete_user_variable(&key, name);
.variables
.delete_user_variable(&ctx.room_id, &ctx.username, name);
let value = match result { let value = match result {
Ok(()) => format!("{} now unset", name), Ok(()) => format!("{} now unset", name),

View File

@ -6,7 +6,7 @@ use thiserror::Error;
/// Shortcut to defining db migration versions. Will probably /// Shortcut to defining db migration versions. Will probably
/// eventually be moved to a config file. /// eventually be moved to a config file.
const MIGRATION_VERSION: u32 = 4; const MIGRATION_VERSION: u32 = 5;
#[derive(Error, Debug)] #[derive(Error, Debug)]
pub enum ConfigError { pub enum ConfigError {

View File

@ -3,20 +3,25 @@ use crate::db::Database;
/// A context carried through the system providing access to things /// A context carried through the system providing access to things
/// like the database. /// like the database.
#[derive(Clone)] #[derive(Clone)]
pub struct Context { pub struct Context<'a> {
pub db: Database, pub db: Database,
pub room_id: String, pub room_id: &'a str,
pub username: String, pub username: &'a str,
pub message_body: String, pub message_body: &'a str,
} }
impl Context { impl<'a> Context<'a> {
pub fn new(db: &Database, room_id: &str, username: &str, message_body: &str) -> Context { pub fn new(
db: &Database,
room_id: &'a str,
username: &'a str,
message_body: &'a str,
) -> Context<'a> {
Context { Context {
db: db.clone(), db: db.clone(),
room_id: room_id.to_owned(), room_id: room_id,
username: username.to_owned(), username: username,
message_body: message_body.to_owned(), message_body: message_body,
} }
} }
} }

View File

@ -2,7 +2,7 @@ use crate::db::errors::{DataError, MigrationError};
use crate::db::migrations::{get_migration_version, Migrations}; use crate::db::migrations::{get_migration_version, Migrations};
use crate::db::variables::Variables; use crate::db::variables::Variables;
use log::info; use log::info;
use sled::Db; use sled::{Config, Db};
use std::path::Path; use std::path::Path;
pub mod data_migrations; pub mod data_migrations;
@ -19,18 +19,27 @@ pub struct Database {
} }
impl Database { impl Database {
pub fn new<P: AsRef<Path>>(path: P) -> Result<Database, DataError> { fn new_db(db: sled::Db) -> Result<Database, DataError> {
let db = sled::open(path)?;
let variables = db.open_tree("variables")?;
let migrations = db.open_tree("migrations")?; let migrations = db.open_tree("migrations")?;
Ok(Database { Ok(Database {
db: db.clone(), db: db.clone(),
variables: Variables(variables), variables: Variables::new(&db)?,
migrations: Migrations(migrations), migrations: Migrations(migrations),
}) })
} }
pub fn new<P: AsRef<Path>>(path: P) -> Result<Database, DataError> {
let db = sled::open(path)?;
Self::new_db(db)
}
pub fn new_temp() -> Result<Database, DataError> {
let config = Config::new().temporary(true);
let db = config.open()?;
Self::new_db(db)
}
pub fn migrate(&self, to_version: u32) -> Result<(), DataError> { pub fn migrate(&self, to_version: u32) -> Result<(), DataError> {
//get version from db //get version from db
let db_version = get_migration_version(&self)?; let db_version = get_migration_version(&self)?;

View File

@ -9,7 +9,9 @@ static MIGRATIONS: phf::Map<u32, DataMigration> = phf_map! {
1u32 => (add_room_user_variable_count::migrate, "add_room_user_variable_count"), 1u32 => (add_room_user_variable_count::migrate, "add_room_user_variable_count"),
2u32 => (delete_v0_schema, "delete_v0_schema"), 2u32 => (delete_v0_schema, "delete_v0_schema"),
3u32 => (delete_variable_count, "delete_variable_count"), 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<Vec<DataMigration>, MigrationError> { pub fn get_migrations(versions: &[u32]) -> Result<Vec<DataMigration>, MigrationError> {

View File

@ -20,6 +20,9 @@ pub enum DataError {
#[error("value does not exist for key: {0}")] #[error("value does not exist for key: {0}")]
KeyDoesNotExist(String), KeyDoesNotExist(String),
#[error("too many entries")]
TooManyEntries,
#[error("expected i32, but i32 schema was violated")] #[error("expected i32, but i32 schema was violated")]
I32SchemaViolation, I32SchemaViolation,
@ -37,6 +40,9 @@ pub enum DataError {
#[error("data migration error: {0}")] #[error("data migration error: {0}")]
MigrationError(#[from] MigrationError), MigrationError(#[from] MigrationError),
#[error("deserialization error: {0}")]
DeserializationError(#[from] bincode::Error),
} }
/// This From implementation is necessary to deal with the recursive /// This From implementation is necessary to deal with the recursive

View File

@ -2,6 +2,7 @@ use crate::db::errors::DataError;
use crate::db::schema::convert_i32; use crate::db::schema::convert_i32;
use byteorder::LittleEndian; use byteorder::LittleEndian;
use sled::transaction::{abort, TransactionalTree}; use sled::transaction::{abort, TransactionalTree};
use sled::Transactional;
use sled::Tree; use sled::Tree;
use std::collections::HashMap; use std::collections::HashMap;
use std::convert::From; use std::convert::From;
@ -11,74 +12,48 @@ use zerocopy::AsBytes;
pub(super) mod migrations; 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)] #[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<D: Into<Vec<u8>>>(space: &[u8], delineator: D) -> Vec<u8> { /// Request soemthing by a username and room ID.
let mut metadata_prefix = vec![]; pub struct UserAndRoom<'a>(pub &'a str, pub &'a str);
metadata_prefix.extend_from_slice(space);
metadata_prefix.push(0xff);
let delineator = delineator.into();
if delineator.len() > 0 { fn to_vec(value: &UserAndRoom<'_>) -> Vec<u8> {
metadata_prefix.extend_from_slice(delineator.as_bytes()); let mut bytes = vec![];
metadata_prefix.push(0xff); 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<u8> {
fn from(value: UserAndRoom) -> Vec<u8> {
to_vec(&value)
} }
metadata_prefix
} }
fn metadata_space_prefix<D: Into<Vec<u8>>>(delineator: D) -> Vec<u8> { impl From<&UserAndRoom<'_>> for Vec<u8> {
space_prefix(METADATA_SPACE, delineator) fn from(value: &UserAndRoom) -> Vec<u8> {
} to_vec(value)
fn metadata_space_key<D: Into<Vec<u8>>>(delineator: D, key_name: &str) -> Vec<u8> {
let mut metadata_key = metadata_space_prefix(delineator);
metadata_key.extend_from_slice(key_name.as_bytes());
metadata_key
}
fn variables_space_prefix<D: Into<Vec<u8>>>(delineator: D) -> Vec<u8> {
space_prefix(VARIABLE_SPACE, delineator)
}
fn variables_space_key<D: Into<Vec<u8>>>(delineator: D, key_name: &str) -> Vec<u8> {
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<RoomAndUser<'a>> for Vec<u8> {
fn from(value: RoomAndUser<'a>) -> Vec<u8> {
let mut bytes = vec![];
bytes.extend_from_slice(value.0.as_bytes());
bytes.push(0xfe);
bytes.extend_from_slice(value.1.as_bytes());
bytes
} }
} }
/// Use a transaction to atomically alter the count of variables in /// Use a transaction to atomically alter the count of variables in
/// the database by the given amount. Count cannot go below 0. /// the database by the given amount. Count cannot go below 0.
fn alter_room_variable_count( fn alter_room_variable_count(
variables: &TransactionalTree, room_variable_count: &TransactionalTree,
room_id: &str, user_and_room: &UserAndRoom<'_>,
username: &str,
amount: i32, amount: i32,
) -> Result<i32, DataError> { ) -> Result<i32, DataError> {
let key = metadata_space_key(RoomAndUser(room_id, username), VARIABLE_COUNT_KEY); let key: Vec<u8> = 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, Some(bytes) => convert_i32(&bytes)? + amount,
None => amount, None => amount,
}; };
@ -88,21 +63,28 @@ fn alter_room_variable_count(
} }
let db_value: I32<LittleEndian> = I32::new(new_count); let db_value: I32<LittleEndian> = I32::new(new_count);
variables.insert(key, db_value.as_bytes())?; room_variable_count.insert(key, db_value.as_bytes())?;
Ok(new_count) Ok(new_count)
} }
impl Variables { impl Variables {
pub(in crate::db) fn new(db: &sled::Db) -> Result<Variables, sled::Error> {
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( pub fn get_user_variables(
&self, &self,
room_id: &str, key: &UserAndRoom<'_>,
username: &str,
) -> Result<HashMap<String, i32>, DataError> { ) -> Result<HashMap<String, i32>, DataError> {
let prefix = variables_space_prefix(RoomAndUser(room_id, username)); let mut prefix: Vec<u8> = key.into();
prefix.push(0xff);
let prefix_len: usize = prefix.len(); let prefix_len: usize = prefix.len();
let variables: Result<Vec<_>, DataError> = self let variables: Result<Vec<_>, DataError> = self
.0 .room_user_variables
.scan_prefix(prefix) .scan_prefix(prefix)
.map(|entry| match entry { .map(|entry| match entry {
Ok((key, raw_value)) => { Ok((key, raw_value)) => {
@ -114,99 +96,104 @@ impl Variables {
}) })
.collect(); .collect();
//Convert I32 to hash map. Can we do this in the first mapping //Convert I32 to hash map. collect() inferred via return type.
//step instead? For some reason this is faster.
variables.map(|entries| entries.into_iter().collect()) variables.map(|entries| entries.into_iter().collect())
} }
pub fn get_variable_count(&self, room_id: &str, username: &str) -> Result<i32, DataError> { pub fn get_variable_count(&self, user_and_room: &UserAndRoom<'_>) -> Result<i32, DataError> {
let delineator = RoomAndUser(room_id, username); let key: Vec<u8> = user_and_room.into();
let key = metadata_space_key(delineator, VARIABLE_COUNT_KEY);
if let Some(raw_value) = self.0.get(&key)? { match self.room_user_variable_count.get(&key)? {
convert_i32(&raw_value) Some(raw_value) => convert_i32(&raw_value),
} else { None => Ok(0),
Ok(0)
} }
} }
pub fn get_user_variable( pub fn get_user_variable(
&self, &self,
room_id: &str, user_and_room: &UserAndRoom<'_>,
username: &str,
variable_name: &str, variable_name: &str,
) -> Result<i32, DataError> { ) -> Result<i32, DataError> {
let key = variables_space_key(RoomAndUser(room_id, username), variable_name); let mut key: Vec<u8> = user_and_room.into();
key.push(0xff);
key.extend_from_slice(variable_name.as_bytes());
if let Some(raw_value) = self.0.get(&key)? { match self.room_user_variables.get(&key)? {
convert_i32(&raw_value) Some(raw_value) => convert_i32(&raw_value),
} else { _ => Err(DataError::KeyDoesNotExist(variable_name.to_owned())),
Err(DataError::KeyDoesNotExist(variable_name.to_owned()))
} }
} }
pub fn set_user_variable( pub fn set_user_variable(
&self, &self,
room_id: &str, user_and_room: &UserAndRoom<'_>,
username: &str,
variable_name: &str, variable_name: &str,
value: i32, value: i32,
) -> Result<(), DataError> { ) -> Result<(), DataError> {
self.0 if self.get_variable_count(user_and_room)? >= 100 {
.transaction(|tx| { return Err(DataError::TooManyEntries);
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<u8> = user_and_room.into();
key.push(0xff);
key.extend_from_slice(variable_name.as_bytes());
let db_value: I32<LittleEndian> = I32::new(value); let db_value: I32<LittleEndian> = 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. //Only increment variable count on new keys.
if let None = old_value { if let None = old_value {
match alter_room_variable_count(&tx, room_id, username, 1) { if let Err(e) = alter_room_variable_count(&tx_counts, &user_and_room, 1) {
Err(e) => abort(e), return abort(e);
_ => Ok(()),
} }
} else {
Ok(())
} }
})
.map_err(|e| e.into()) Ok(())
},
)?;
Ok(())
} }
pub fn delete_user_variable( pub fn delete_user_variable(
&self, &self,
room_id: &str, user_and_room: &UserAndRoom<'_>,
username: &str,
variable_name: &str, variable_name: &str,
) -> Result<(), DataError> { ) -> Result<(), DataError> {
self.0 (&self.room_user_variables, &self.room_user_variable_count).transaction(
.transaction(|tx| { |(tx_vars, tx_counts)| {
let key = variables_space_key(RoomAndUser(room_id, username), variable_name); let mut key: Vec<u8> = user_and_room.into();
key.push(0xff);
key.extend_from_slice(variable_name.as_bytes());
//TODO why does tx.remove require moving the key? //TODO why does tx.remove require moving the key?
if let Some(_) = tx.remove(key.clone())? { if let Some(_) = tx_vars.remove(key.clone())? {
match alter_room_variable_count(&tx, room_id, username, -1) { if let Err(e) = alter_room_variable_count(&tx_counts, user_and_room, -1) {
Err(e) => abort(e), return abort(e);
_ => Ok(()),
} }
} else { } 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::tempdir; use sled::Config;
fn create_test_instance() -> Variables { fn create_test_instance() -> Variables {
Variables( let config = Config::new().temporary(true);
sled::open(&tempdir().unwrap()) let db = config.open().unwrap();
.unwrap() Variables::new(&db).unwrap()
.open_tree("variables")
.unwrap(),
)
} }
//Room Variable count tests //Room Variable count tests
@ -214,24 +201,23 @@ mod tests {
#[test] #[test]
fn alter_room_variable_count_test() { fn alter_room_variable_count_test() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
let alter_count = |amount: i32| { let alter_count = |amount: i32| {
variables variables
.0 .room_user_variable_count
.transaction(|tx| { .transaction(|tx| match alter_room_variable_count(&tx, &key, amount) {
match alter_room_variable_count(&tx, "room", "username", amount) { Err(e) => abort(e),
Err(e) => abort(e), _ => Ok(()),
_ => Ok(()),
}
}) })
.expect("got transaction failure"); .expect("got transaction failure");
}; };
fn get_count(variables: &Variables) -> i32 { let get_count = |variables: &Variables| -> i32 {
variables variables
.get_variable_count("room", "username") .get_variable_count(&key)
.expect("could not get variable count") .expect("could not get variable count")
} };
//addition //addition
alter_count(5); alter_count(5);
@ -245,19 +231,18 @@ mod tests {
#[test] #[test]
fn alter_room_variable_count_cannot_go_below_0_test() { fn alter_room_variable_count_cannot_go_below_0_test() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
variables variables
.0 .room_user_variable_count
.transaction( .transaction(|tx| match alter_room_variable_count(&tx, &key, -1000) {
|tx| match alter_room_variable_count(&tx, "room", "username", -1000) { Err(e) => abort(e),
Err(e) => abort(e), _ => Ok(()),
_ => Ok(()), })
},
)
.expect("got transaction failure"); .expect("got transaction failure");
let count = variables let count = variables
.get_variable_count("room", "username") .get_variable_count(&key)
.expect("could not get variable count"); .expect("could not get variable count");
assert_eq!(0, count); assert_eq!(0, count);
@ -266,9 +251,10 @@ mod tests {
#[test] #[test]
fn empty_db_reports_0_room_variable_count_test() { fn empty_db_reports_0_room_variable_count_test() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
let count = variables let count = variables
.get_variable_count("room", "username") .get_variable_count(&key)
.expect("could not get variable count"); .expect("could not get variable count");
assert_eq!(0, count); assert_eq!(0, count);
@ -277,13 +263,14 @@ mod tests {
#[test] #[test]
fn set_user_variable_increments_count() { fn set_user_variable_increments_count() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
variables variables
.set_user_variable("room", "username", "myvariable", 5) .set_user_variable(&key, "myvariable", 5)
.expect("could not insert variable"); .expect("could not insert variable");
let count = variables let count = variables
.get_variable_count("room", "username") .get_variable_count(&key)
.expect("could not get variable count"); .expect("could not get variable count");
assert_eq!(1, count); assert_eq!(1, count);
@ -292,17 +279,18 @@ mod tests {
#[test] #[test]
fn update_user_variable_does_not_increment_count() { fn update_user_variable_does_not_increment_count() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
variables variables
.set_user_variable("room", "username", "myvariable", 5) .set_user_variable(&key, "myvariable", 5)
.expect("could not insert variable"); .expect("could not insert variable");
variables variables
.set_user_variable("room", "username", "myvariable", 10) .set_user_variable(&key, "myvariable", 10)
.expect("could not update variable"); .expect("could not update variable");
let count = variables let count = variables
.get_variable_count("room", "username") .get_variable_count(&key)
.expect("could not get variable count"); .expect("could not get variable count");
assert_eq!(1, count); assert_eq!(1, count);
@ -313,30 +301,49 @@ mod tests {
#[test] #[test]
fn set_and_get_variable_test() { fn set_and_get_variable_test() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
variables variables
.set_user_variable("room", "username", "myvariable", 5) .set_user_variable(&key, "myvariable", 5)
.expect("could not insert variable"); .expect("could not insert variable");
let value = variables let value = variables
.get_user_variable("room", "username", "myvariable") .get_user_variable(&key, "myvariable")
.expect("could not get value"); .expect("could not get value");
assert_eq!(5, 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] #[test]
fn delete_variable_test() { fn delete_variable_test() {
let variables = create_test_instance(); let variables = create_test_instance();
let key = UserAndRoom("username", "room");
variables variables
.set_user_variable("room", "username", "myvariable", 5) .set_user_variable(&key, "myvariable", 5)
.expect("could not insert variable"); .expect("could not insert variable");
variables variables
.delete_user_variable("room", "username", "myvariable") .delete_user_variable(&key, "myvariable")
.expect("could not delete value"); .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!(result.is_err());
assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_))));
@ -345,7 +352,8 @@ mod tests {
#[test] #[test]
fn get_missing_variable_returns_key_does_not_exist() { fn get_missing_variable_returns_key_does_not_exist() {
let variables = create_test_instance(); 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!(result.is_err());
assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_))));
@ -354,7 +362,8 @@ mod tests {
#[test] #[test]
fn remove_missing_variable_returns_key_does_not_exist() { fn remove_missing_variable_returns_key_does_not_exist() {
let variables = create_test_instance(); 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!(result.is_err());
assert!(matches!(result, Err(DataError::KeyDoesNotExist(_)))); assert!(matches!(result, Err(DataError::KeyDoesNotExist(_))));

View File

@ -5,7 +5,8 @@ use byteorder::LittleEndian;
use memmem::{Searcher, TwoWaySearcher}; use memmem::{Searcher, TwoWaySearcher};
use sled::transaction::TransactionError; use sled::transaction::TransactionError;
use sled::{Batch, IVec}; use sled::{Batch, IVec};
use zerocopy::byteorder::U32; use std::collections::HashMap;
use zerocopy::byteorder::{I32, U32};
use zerocopy::AsBytes; use zerocopy::AsBytes;
pub(in crate::db) mod add_room_user_variable_count { 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<u8> {
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> { pub(in crate::db) fn migrate(db: &Database) -> Result<(), DataError> {
let tree = &db.variables.0; let tree = &db.variables.room_user_variables;
let prefix = variables_space_prefix(""); let prefix = b"variables";
//Extract a vec of tuples, consisting of room id + username. //Extract a vec of tuples, consisting of room id + username.
let results: Vec<RoomAndUser> = tree let results: Vec<RoomAndUser> = tree
@ -73,13 +85,12 @@ pub(in crate::db) mod add_room_user_variable_count {
//Start a transaction on the variables tree. //Start a transaction on the variables tree.
let tx_result: Result<_, TransactionError<DataError>> = let tx_result: Result<_, TransactionError<DataError>> =
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 batch = counts.iter().fold(Batch::default(), |mut batch, entry| {
let (info, count) = entry; let (info, count) = entry;
//Add variable count according to new schema. //Add variable count according to new schema.
let delineator = super::RoomAndUser(&info.room_id, &info.username); let key = create_key(&info.room_id, &info.username);
let key = variables_space_key(delineator, VARIABLE_COUNT_KEY);
let db_value: U32<LittleEndian> = U32::new(*count); let db_value: U32<LittleEndian> = U32::new(*count);
batch.insert(key, db_value.as_bytes()); 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> { 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(); let mut batch = Batch::default();
while let Some(Ok((key, _))) = vars.next() { 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(()) Ok(())
} }
pub(in crate::db) fn delete_variable_count(db: &Database) -> Result<(), DataError> { pub(in crate::db) fn delete_variable_count(db: &Database) -> Result<(), DataError> {
let prefix = variables_space_prefix(""); let prefix = b"variables";
let mut vars = db.variables.0.scan_prefix(prefix); let mut vars = db.variables.room_user_variables.scan_prefix(prefix);
let mut batch = Batch::default(); let mut batch = Batch::default();
while let Some(Ok((key, _))) = vars.next() { 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(()) Ok(())
} }
@ -203,8 +214,8 @@ pub(in crate::db) mod change_delineator_delimiter {
} }
pub fn migrate(db: &Database) -> Result<(), DataError> { pub fn migrate(db: &Database) -> Result<(), DataError> {
let tree = &db.variables.0; let tree = &db.variables.room_user_variables;
let prefix = variables_space_prefix(""); let prefix = b"variables";
let results: Vec<UserVariableEntry> = tree let results: Vec<UserVariableEntry> = tree
.scan_prefix(&prefix) .scan_prefix(&prefix)
@ -214,8 +225,8 @@ pub(in crate::db) mod change_delineator_delimiter {
let mut batch = Batch::default(); let mut batch = Batch::default();
for insert in results { for insert in results {
let old = create_old_key(&prefix, &insert); let old = create_old_key(prefix, &insert);
let new = create_new_key(&prefix, &insert); let new = create_new_key(prefix, &insert);
batch.remove(old); batch.remove(old);
batch.insert(new, insert.value); batch.insert(new, insert.value);
@ -225,3 +236,119 @@ pub(in crate::db) mod change_delineator_delimiter {
Ok(()) 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<UserVariableEntry, MigrationError> {
if let Ok((key, value)) = entry {
let keys: Vec<Result<&str, _>> = 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<u8> {
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<u8> {
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<UserVariableEntry> = variables_tree
.scan_prefix(&prefix)
.map(extract_v1_entries)
.collect::<Result<Vec<_>, 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<LittleEndian> = I32::new(count);
count_batch.insert(key, db_value.as_bytes());
});
variables_tree.apply_batch(batch)?;
count_tree.apply_batch(count_batch)?;
Ok(())
}
}