Merge pull request #7 from rex-rs/membench

 - tools/memcached_benchmark: refactor into multiple modules

Signed-off-by: Jinghao Jia <jinghao7@illinois.edu>

Author: jinghao-jia

tools/memcached_benchmark/src/cli.rs

@@ -0,0 +1,91 @@
+use clap::{Parser, Subcommand};
+
+use crate::Protocol;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+pub(crate) struct Cli {
+    #[command(subcommand)]
+    pub(crate) command: Commands,
+}
+
+#[derive(Debug, Subcommand)]
+pub(crate) enum Commands {
+    #[command(arg_required_else_help = true)]
+    Bench {
+        /// memcached server addr
+        #[arg(short, long, required = true)]
+        server_address: String,
+
+        #[arg(short, long, default_value = "11211")]
+        port: String,
+
+        /// key size to generate random memcached key
+        #[arg(short, long, default_value = "16")]
+        key_size: usize,
+
+        /// value size to generate random memcached value
+        #[arg(short, long, default_value = "32")]
+        value_size: usize,
+
+        /// verify the value after get command
+        #[arg(long, default_value = "false")]
+        validate: bool,
+
+        /// number of test entries to generate
+        #[arg(short, long, default_value = "100000")]
+        nums: usize,
+
+        /// number of threads to run
+        #[arg(short, long, default_value = "4")]
+        threads: usize,
+
+        /// udp or tcp protocol for memcached
+        #[arg(short = 'l', long, default_value_t = Protocol::Udp , value_enum)]
+        protocol: Protocol,
+
+        /// number of dict entries to generate
+        #[arg(short, long, default_value = "1000000")]
+        dict_entries: usize,
+
+        /// load the prepared test_entries from disk
+        #[arg(long, default_value = "false")]
+        load_bench_entries: bool,
+
+        /// skip set memcached value if the data is already imported
+        #[arg(long, default_value = "false")]
+        skip_set: bool,
+
+        /// bounded mpsc channel for communicating between asynchronous tasks
+        /// with backpressure
+        #[arg(long, default_value = "200")]
+        pipeline: usize,
+
+        /// dict path to load
+        #[arg(
+            short = 'f',
+            long,
+            default_value = "test_dict.yml.zst",
+            conflicts_with = "key_size",
+            conflicts_with = "dict_entries"
+        )]
+        dict_path: String,
+    },
+    GenTestdict {
+        /// key size to generate random memcached key
+        #[arg(short, long, default_value = "16")]
+        key_size: usize,
+
+        /// value size to generate random memcached value
+        #[arg(short, long, default_value = "32")]
+        value_size: usize,
+
+        /// number of dict entries to generate
+        #[arg(short, long, default_value = "1000000")]
+        dict_entries: usize,
+
+        /// dict path to store
+        #[arg(short = 'f', long, default_value = "test_dict.yml.zst")]
+        dict_path: String,
+    },
+}

tools/memcached_benchmark/src/dict.rs

@@ -0,0 +1,128 @@
+use std::{collections::HashMap, mem::size_of_val, sync::Arc};
+
+use log::{debug, info};
+use rand::{
+    Rng,
+    distr::{Alphanumeric, SampleString},
+};
+use rand_chacha::{ChaCha8Rng, rand_core::SeedableRng};
+use rand_distr::Zipf;
+use rayon::prelude::*;
+
+use crate::{Protocol, fs::write_hashmap_to_file};
+
+const SEED: u64 = 12312;
+
+/// Generate random string of specified length
+pub(crate) fn generate_random_str(len: usize) -> String {
+    Alphanumeric.sample_string(&mut rand::rng(), len)
+}
+
+/// Generate test dictionary for memcached with random keys and values
+pub(crate) fn generate_memcached_test_dict(
+    key_size: usize,
+    value_size: usize,
+    nums: usize,
+) -> HashMap<String, String> {
+    // random generate dict for memcached test
+    (0..nums)
+        .into_par_iter()
+        .map(|_| {
+            (
+                generate_random_str(key_size),
+                generate_random_str(value_size),
+            )
+        })
+        .collect()
+}
+
+/// Generate test dict and write to disk
+/// # Arguments
+/// * `key_size` - key size
+/// * `value_size` - value size
+/// * `nums` - number of entries
+/// * `dict_path` - dict path to store
+/// # Returns
+/// * `Result` - Result<HashMap<String, String>, anyhow::Error>
+/// # Example
+/// ```rust
+/// let test_dict = generate_test_dict_write_to_disk(16, 32, 100000, "test_dict.yml.zst");
+/// ```
+pub(crate) fn generate_test_dict_write_to_disk(
+    key_size: usize,
+    value_size: usize,
+    nums: usize,
+    dict_path: &str,
+) -> anyhow::Result<HashMap<String, String>> {
+    let test_dict = generate_memcached_test_dict(key_size, value_size, nums);
+    debug!("test dict len: {}", test_dict.len());
+    if let Some((key, value)) = test_dict.iter().next() {
+        debug!("test dict key size: {}", size_of_val(key.as_str()));
+        debug!("test dict value size: {}", size_of_val(value.as_str()));
+    } else {
+        return Err(anyhow::anyhow!("test dict is empty"));
+    }
+    write_hashmap_to_file(&test_dict, dict_path)?;
+    info!("write test dict to path {}", dict_path);
+    Ok(test_dict)
+}
+
+/// Generate test entries with Zipf distribution for benchmarking
+pub(crate) fn generate_test_entries(
+    test_dict: Arc<HashMap<Arc<String>, Arc<String>>>,
+    nums: usize,
+) -> Vec<(Arc<String>, Arc<String>, Protocol)> {
+    let mut rng = ChaCha8Rng::seed_from_u64(SEED);
+    let zipf = Zipf::new((test_dict.len() - 1) as f64, 0.99).unwrap();
+
+    let keys: Vec<Arc<String>> = test_dict.keys().cloned().collect();
+    (0..nums)
+        .map(|idx| {
+            let key = &keys[rng.sample(zipf) as usize];
+            let value = test_dict.get(key).unwrap();
+            // every 31 element is tcp. udp:tcp = 30:1
+            let protocol = if idx % 31 == 30 {
+                Protocol::Tcp
+            } else {
+                Protocol::Udp
+            };
+            (key.clone(), value.clone(), protocol)
+        })
+        .collect()
+}
+
+/// Analyze the statistics of test entries
+#[allow(dead_code)]
+pub(crate) fn test_entries_statistics(
+    test_entries: Arc<Vec<(&String, &String, Protocol)>>,
+) {
+    let mut udp_count: usize = 0;
+    let mut tcp_count: usize = 0;
+
+    // analyze the key distribution base on the frequency
+    let mut key_frequency = HashMap::new();
+
+    // only get the first element in the tuple
+    test_entries.iter().for_each(|(key, _, proto)| {
+        *key_frequency.entry(key.to_string()).or_insert(0) += 1;
+        if *proto == Protocol::Udp {
+            udp_count += 1;
+        } else {
+            tcp_count += 1;
+        }
+    });
+
+    // sort by frequency
+    let mut key_frequency: Vec<_> = key_frequency.into_iter().collect();
+    key_frequency.sort_by(|a, b| a.1.cmp(&b.1));
+
+    // Display the frequency of each item
+    for (key, count) in &key_frequency {
+        if *count < key_frequency.len() / 1000 {
+            continue;
+        }
+        info!("{}: {}", key, count);
+    }
+
+    info!("tcp count: {}, udp count: {}", tcp_count, udp_count);
+}

tools/memcached_benchmark/src/fs.rs

@@ -0,0 +1,92 @@
+use std::{
+    collections::HashMap,
+    fs::File,
+    io::{BufRead, BufReader, Write},
+    path::Path,
+};
+
+use anyhow::Result;
+use log::{debug, info};
+use serde::Serialize;
+
+use crate::Protocol;
+
+/// Write hashmap to a compressed file using zstd
+/// # Arguments
+/// * `hashmap` - hashmap to write to file
+/// * `file_path` - file path to write
+/// # Returns
+/// * `Result` - Result<(), anyhow::Error>
+pub(crate) fn write_hashmap_to_file<T: Serialize>(
+    hashmap: &T,
+    file_path: &str,
+) -> Result<()> {
+    // Serialize the hashmap to a YAML string
+    let serialized =
+        serde_yaml::to_string(hashmap).expect("Failed to serialize");
+
+    // Create or open the file
+    let file = File::create(file_path)?;
+
+    // Create a zstd encoder with compression level 7
+    let mut encoder = zstd::stream::write::Encoder::new(file, 7)?;
+
+    // Write the YAML string to the file
+    encoder.write_all(serialized.as_bytes())?;
+    encoder.finish()?;
+
+    Ok(())
+}
+
+/// Loads benchmark entries from disk from a zstd-compressed YAML file.
+///
+/// Although the entries are typically selected randomly from a test dictionary,
+/// there are cases where a fixed sequence of entries is needed to ensure
+/// consistent performance comparisons, and this function is utilized to
+/// retrieve the stored benchmark entries
+pub(crate) fn load_bench_entries_from_disk(
+    path: &Path,
+) -> Vec<(String, String, Protocol)> {
+    let file = std::fs::File::open(path).unwrap();
+    let decoder = zstd::stream::read::Decoder::new(file).unwrap();
+    let reader = std::io::BufReader::new(decoder);
+    let test_entries: Vec<(String, String, Protocol)> =
+        serde_yaml::from_reader(reader).unwrap();
+    test_entries
+}
+
+/// Load test dictionary from disk
+///
+/// This function opens a file located at `test_dict_path`, which is expected to
+/// be a zstd-compressed and valid YAML document key-value pair
+/// (`HashMap<String, String>`).
+pub(crate) fn load_test_dict(
+    test_dict_path: &std::path::Path,
+) -> anyhow::Result<HashMap<String, String>> {
+    // load dict from file if dict_path is not empty
+    info!("loading dict from path {:?}", test_dict_path);
+    let file = std::fs::File::open(test_dict_path)?;
+    let decoder = zstd::stream::read::Decoder::new(file)?;
+    let reader = BufReader::new(decoder);
+
+    // Deserialize the string into a HashMap
+    let mut test_dict = HashMap::new();
+
+    reader.lines().for_each(|line| {
+        let line = line.unwrap();
+        // Assuming each line in your file is a valid YAML representing a
+        // key-value pair
+        let deserialized_map: HashMap<String, String> =
+            serde_yaml::from_str(&line).unwrap();
+        test_dict.extend(deserialized_map);
+    });
+
+    debug!("test dict len: {}", test_dict.len());
+    if let Some(key) = test_dict.keys().next() {
+        debug!("test dict key size: {}", key.len());
+    }
+    if let Some(value) = test_dict.values().next() {
+        debug!("test dict value size: {}", value.len());
+    }
+    Ok(test_dict)
+}