projectmoon
/
tenebrous-dicebot
1
0
Fork 1
Code Issues 12 Pull Requests 1 Projects Releases 10 Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
Cargo.lock generated
View File

@ -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",

6
Cargo.toml
View File

@ -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" ]

2
src/bin/dicebot-cmd.rs
View File

@ -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::<Vec<String>>().join(" ");
let command = match commands::parse(&input) {
Ok(command) => command,

24
src/cofd/dice.rs
View File

@ -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<i32, BotError> {
async fn calculate_dice_amount(pool: &DicePoolWithContext<'_>) -> Result<i32, BotError> {
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<i32, BotError> = 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<RolledDicePool,
mod tests {
use super::*;
use crate::db::Database;
use tempfile::tempdir;
///Instead of being random, generate a series of numbers we have complete
///control over.
@ -507,7 +502,7 @@ mod tests {
#[tokio::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 mut amounts = vec![];
@ -532,7 +527,7 @@ mod tests {
#[tokio::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 mut amounts = vec![];
@ -554,11 +549,12 @@ mod tests {
#[tokio::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 user_and_room = crate::db::variables::UserAndRoom(&ctx.username, &ctx.room_id);
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");
let amounts = vec![Amount {

4
src/commands.rs
View File

@ -37,7 +37,7 @@ impl Execution {
#[async_trait]
pub trait Command: Send + Sync {
async fn execute(&self, ctx: &Context) -> 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<CommandResult> {
pub async fn execute_command(ctx: &Context<'_>) -> Option<CommandResult> {
let res = parse(&ctx.message_body);
let (plain, html) = match res {

2
src/commands/basic_rolling.rs
View File

@ -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!(

2
src/commands/cofd.rs
View File

@ -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;

4
src/commands/cthulhu.rs
View File

@ -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);

2
src/commands/misc.rs
View File

@ -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",

33
src/commands/variables.rs
View File

@ -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),

2
src/config.rs
View File

@ -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 {

23
src/context.rs
View File

@ -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,
}
}
}

19
src/db.rs
View File

@ -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<P: AsRef<Path>>(path: P) -> Result<Database, DataError> {
let db = sled::open(path)?;
let variables = db.open_tree("variables")?;
fn new_db(db: sled::Db) -> Result<Database, DataError> {
let migrations = db.open_tree("migrations")?;
Ok(Database {
db: db.clone(),
variables: Variables(variables),
variables: Variables::new(&db)?,
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> {
//get version from db
let db_version = get_migration_version(&self)?;

4
src/db/data_migrations.rs
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"),
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<Vec<DataMigration>, MigrationError> {

6
src/db/errors.rs
View File

@ -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

279
src/db/variables.rs
View File

@ -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<D: Into<Vec<u8>>>(space: &[u8], delineator: D) -> Vec<u8> {
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<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
}
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> {
space_prefix(METADATA_SPACE, delineator)
}
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
impl From<&UserAndRoom<'_>> for Vec<u8> {
fn from(value: &UserAndRoom) -> Vec<u8> {
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<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,
None => amount,
};
@ -88,21 +63,28 @@ fn alter_room_variable_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)
}
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(
&self,
room_id: &str,
username: &str,
key: &UserAndRoom<'_>,
) -> 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 variables: Result<Vec<_>, 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<i32, DataError> {
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<i32, DataError> {
let key: Vec<u8> = 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<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)? {
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<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 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<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?
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(_))));

157
src/db/variables/migrations.rs
View File

@ -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<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> {
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<RoomAndUser> = 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<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 (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<LittleEndian> = 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<UserVariableEntry> = 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<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(())
}
}