Skip to content

Commit

Permalink
refactor(torture): use OrdMap to track the state
Browse files Browse the repository at this point in the history
  • Loading branch information
gabriele-0201 committed Jan 21, 2025
1 parent be67214 commit 544b136
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 42 deletions.
1 change: 0 additions & 1 deletion torture/src/agent.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use anyhow::{anyhow, bail, Result};
use futures::SinkExt as _;
use nomt::{Blake3Hasher, Nomt};
use rand::Rng;
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::{
io::{BufReader, BufWriter},
Expand Down
7 changes: 7 additions & 0 deletions torture/src/message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,13 @@ impl KeyValueChange {
KeyValueChange::Insert(ref key, _) | KeyValueChange::Delete(ref key) => key,
}
}

pub fn value(&self) -> Option<Value> {
match self {
KeyValueChange::Insert(_, val) => Some(val.clone()),
KeyValueChange::Delete(_) => None,
}
}
}

/// The parameters for the [`ToAgent::Init`] message.
Expand Down
81 changes: 40 additions & 41 deletions torture/src/supervisor/workload.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use anyhow::Result;
use imbl::OrdMap;
use rand::prelude::*;
use std::{collections::BTreeMap, time::Duration};
use std::time::Duration;
use tempfile::TempDir;
use tokio::time::{error::Elapsed, timeout};
use tokio_util::sync::CancellationToken;
Expand Down Expand Up @@ -48,10 +48,7 @@ impl Biases {
struct Snapshot {
sync_seqn: u32,
/// The state of the database at this snapshot.
///
/// The key is the key of the key-value pair. Instead of using the 32 bytes of the key, we
/// save the index of the key in the `WorkloadState::mentions` vector.
state: OrdMap<u32, Option<Vec<u8>>>,
state: OrdMap<[u8; 32], Option<Vec<u8>>>,
}

impl Snapshot {
Expand All @@ -73,11 +70,6 @@ struct WorkloadState {
random_size: bool,
/// The values that were committed.
committed: Snapshot,
/// All keys that were generated during this run.
///
/// It's extremely unlikely to find any duplicates in there. Each key is very likely to store
/// a value.
mentions: Vec<[u8; 32]>,
}

impl WorkloadState {
Expand All @@ -88,7 +80,6 @@ impl WorkloadState {
size,
random_size,
committed: Snapshot::empty(),
mentions: Vec::with_capacity(32 * 1024),
}
}

Expand All @@ -112,47 +103,55 @@ impl WorkloadState {
// Commiting requires using only the unique keys. To ensure that we deduplicate the keys
// using a hash set.
let mut used_keys = std::collections::HashSet::with_capacity(num_changes);
for _ in 0..num_changes {
// Generate a key until unique.
let key_index = loop {
let ki = self.gen_key();
if used_keys.insert(ki) {
break ki;
}
};

if self.rng.gen_bool(self.biases.delete) {
snapshot.state.insert(key_index, None);
let key = self.mentions[key_index as usize];
changes.push(KeyValueChange::Delete(key));
} else {
let value = self.gen_value();
snapshot.state.insert(key_index, Some(value.clone()));
let key = self.mentions[key_index as usize];
changes.push(KeyValueChange::Insert(key, value));
};
loop {
let change = self.gen_key_value_change();
if used_keys.contains(change.key()) {
continue;
}

snapshot.state.insert(*change.key(), change.value());
used_keys.insert(*change.key());
changes.push(change);

if used_keys.len() >= num_changes {
break;
}
}

changes.sort_by(|a, b| a.key().cmp(&b.key()));

(snapshot, changes)
}

/// Returns a key that was generated before or a new key. The returned value is an index in the
/// `mentions` vector.
fn gen_key(&mut self) -> u32 {
/// Returns a KeyValueChange with a new key, a deleted or a modified one.
fn gen_key_value_change(&mut self) -> KeyValueChange {
// TODO: sophisticated key generation.
//
// - Pick a key that was already generated before, but generate a key that shares some bits.
if self.mentions.is_empty() || self.rng.gen_bool(self.biases.new_key) {
let mut key = [0; 32];
let mut key = [0; 32];
if !self.committed.state.is_empty() && self.rng.gen_bool(self.biases.delete) {
loop {
self.rng.fill_bytes(&mut key);
if let Some((next_key, Some(_))) = self.committed.state.get_next(&key) {
return KeyValueChange::Delete(*next_key);
}
}
}

if self.committed.state.is_empty() || self.rng.gen_bool(self.biases.new_key) {
loop {
self.rng.fill_bytes(&mut key);
if !self.committed.state.contains_key(&key) {
return KeyValueChange::Insert(key, self.gen_value());
}
}
}

loop {
self.rng.fill_bytes(&mut key);
let ix = self.mentions.len() as u32;
self.mentions.push(key.clone());
ix
} else {
let ix = self.rng.gen_range(0..self.mentions.len());
ix as u32
if let Some((next_key, _)) = self.committed.state.get_next(&key) {
return KeyValueChange::Insert(*next_key, self.gen_value());
}
}
}

Expand Down

0 comments on commit 544b136

Please sign in to comment.