feat: support momento leaderboards (#355)

Adds a leaderboard workload type, configuration, metrics, and momento
client implementation.

Author: malandis

Cargo.lock

@@ -35,7 +35,7 @@ hyper-util = { version = "0.1.8", features = ["full"] }
 metriken = "0.7.0"
 metriken-exposition = { version = "0.11.1", features = ["json", "parquet-conversion"] }
 mio = "1.0.3"
-momento = "0.46.1"
+momento = "0.49.0"
 once_cell = "1.18.0"
 openssl = { version = "0.10.71", optional = true }
 openssl-src = "300.3.1"

Cargo.toml

@@ -0,0 +1,80 @@
+# An example configuration for benchmarking Momento (https://www.gomomento.com)
+# leaderboards.
+
+[general]
+# specify the protocol to be used
+protocol = "momento"
+# the interval for stats integration and reporting
+interval = 60
+# the number of intervals to run the test for
+duration = 300
+# run the admin thread with a HTTP listener at the address provided, this allows
+# stats exposition via HTTP
+admin = "127.0.0.1:9090"
+# optionally, set an initial seed for the PRNGs used to generate the workload.
+# The default is to intialize from the OS entropy pool.
+#initial_seed = "0"
+
+#[metrics]
+# output file for detailed stats during the run
+#output = "stats.json"
+# format of the output file (possible values are json, msgpack, parquet)
+#format = "json"
+# optionally specify batch size for parquet row groups
+# only valid for parquet output
+#batch_size = 100_000
+# optionally specify histogram type (can be standard (default) or sparse)
+# only valid for parquet output
+#histogram = "sparse"
+# optionally, specify the sampling interval for metrics. Input is a string
+# with the unit attached; for example "100ms" or "1s". Defaults to 1s.
+#interval = "1s"
+
+[debug]
+# choose from: error, warn, info, debug, trace
+log_level = "info"
+# optionally, log to the file below instead of standard out
+# log_file = "rpc-perf.log"
+# backup file name for use with log rotation
+log_backup = "rpc-perf.log.old"
+# trigger log rotation when the file grows beyond this size (in bytes). Set this
+# option to '0' to disable log rotation.
+log_max_size = 1073741824
+
+[target]
+# we don't need to specify any endpoints for momento
+endpoints = []
+# specify the name of the target cache for the leaderboard
+cache_name = "test-cache"
+
+[leaderboard]
+# number of threads used to drive client requests
+threads = 1
+# number of gRPC clients to initialize, each maintains at least one TCP stream
+poolsize = 1
+# an upper limit on the number of concurrent requests per gRPC client
+concurrency = 20
+# set the timeout in milliseconds
+request_timeout = 1000
+
+[workload]
+# the number of threads that will be used to generate the workload
+threads = 1
+
+[workload.ratelimit]
+# set a global ratelimit for the workload
+start = 50
+
+[[workload.leaderboard]]
+weight = 1
+# Number of distinct leaderboards to create
+nleaderboards = 1
+# Number of distinct ids to use across all leaderboards
+nids = 100
+
+# run upsert vs get_competition_rank in a 10:1 ratio
+# each time, only upsert/query a single id
+commands = [
+	{ verb = "upsert", weight = 10, cardinality = 1 },
+	{ verb = "get_competition_rank", weight = 1, cardinality = 1 },
+]

configs/momento_leaderboard.toml

@@ -0,0 +1,45 @@
+/// # Leaderboard Clients
+///
+/// Leaderboards provide at-scale competitor ranking and lookup.
+///
+/// RPC-Perf store clients are used to evaluate the performance of object
+/// leaderboard services in terms of throughput and latency.
+use crate::*;
+
+use async_channel::Receiver;
+use tokio::runtime::Runtime;
+use workload::{ClientWorkItemKind, LeaderboardClientRequest};
+
+mod momento;
+
+pub fn launch(
+    config: &Config,
+    work_receiver: Receiver<ClientWorkItemKind<LeaderboardClientRequest>>,
+) -> Option<Runtime> {
+    if config.leaderboard().is_none() {
+        debug!("No leaderboard configuration specified");
+        return None;
+    }
+    debug!("Launching clients...");
+
+    config.leaderboard()?;
+
+    // spawn the request drivers on their own runtime
+    let mut client_rt = Builder::new_multi_thread()
+        .enable_all()
+        .worker_threads(config.leaderboard().unwrap().threads())
+        .build()
+        .expect("failed to initialize tokio runtime");
+
+    match config.general().protocol() {
+        Protocol::Momento => momento::launch_tasks(&mut client_rt, config.clone(), work_receiver),
+        protocol => {
+            eprintln!(
+                "leaderboard commands are not supported for the {:?} protocol",
+                protocol
+            );
+        }
+    }
+
+    Some(client_rt)
+}

src/clients/leaderboard/mod.rs

@@ -0,0 +1,183 @@
+use super::record_result;
+use crate::clients::ResponseError;
+use crate::config::Config;
+use crate::metrics::*;
+use crate::workload::{ClientWorkItemKind, LeaderboardClientRequest};
+use crate::{workload, RUNNING};
+use paste::paste;
+
+use async_channel::Receiver;
+use momento::leaderboard::{configurations, LeaderboardClient};
+use momento::CredentialProvider;
+use ringlog::debug;
+use tokio::runtime::Runtime;
+use tokio::time::timeout;
+
+use std::io::{Error, Result};
+use std::sync::atomic::Ordering;
+use std::time::Instant;
+
+/// Launch tasks with one channel per task as gRPC is mux-enabled.
+pub fn launch_tasks(
+    runtime: &mut Runtime,
+    config: Config,
+    work_receiver: Receiver<ClientWorkItemKind<LeaderboardClientRequest>>,
+) {
+    debug!("launching momento protocol tasks");
+
+    let cache_name = config
+        .target()
+        .cache_name()
+        .unwrap_or_else(|| {
+            eprintln!("cache name is not specified in the `target` section");
+            std::process::exit(1);
+        })
+        .to_string();
+
+    for _ in 0..config.leaderboard().unwrap().poolsize() {
+        let client = {
+            let _guard = runtime.enter();
+
+            // initialize the Momento cache client
+            if std::env::var("MOMENTO_API_KEY").is_err() {
+                eprintln!("environment variable `MOMENTO_API_KEY` is not set");
+                std::process::exit(1);
+            }
+
+            let credential_provider =
+                match CredentialProvider::from_env_var("MOMENTO_API_KEY".to_string()) {
+                    Ok(v) => v,
+                    Err(e) => {
+                        eprintln!("MOMENTO_API_KEY key should be valid: {e}");
+                        std::process::exit(1);
+                    }
+                };
+
+            match LeaderboardClient::builder()
+                .configuration(configurations::LowLatency::v1())
+                .credential_provider(credential_provider)
+                .build()
+            {
+                Ok(c) => c,
+                Err(e) => {
+                    eprintln!("could not create leaderboard client: {}", e);
+                    std::process::exit(1);
+                }
+            }
+        };
+
+        LEADERBOARD_CONNECT.increment();
+        LEADERBOARD_CONNECT_CURR.increment();
+
+        // create one task per channel
+        for _ in 0..config.leaderboard().unwrap().concurrency() {
+            runtime.spawn(task(
+                config.clone(),
+                client.clone(),
+                cache_name.clone(),
+                work_receiver.clone(),
+            ));
+        }
+    }
+}
+
+async fn task(
+    config: Config,
+    mut client: LeaderboardClient,
+    cache_name: String,
+    work_receiver: Receiver<ClientWorkItemKind<LeaderboardClientRequest>>,
+) -> Result<()> {
+    while RUNNING.load(Ordering::Relaxed) {
+        let work_item = work_receiver
+            .recv()
+            .await
+            .map_err(|_| Error::other("channel closed"))?;
+
+        LEADERBOARD_REQUEST.increment();
+        let start = Instant::now();
+        let result = match work_item {
+            ClientWorkItemKind::Request { request, .. } => match request {
+                LeaderboardClientRequest::Upsert(r) => {
+                    upsert(&mut client, &config, cache_name.clone(), r).await
+                }
+                LeaderboardClientRequest::GetCompetitionRank(r) => {
+                    get_competition_rank(&mut client, &config, cache_name.clone(), r).await
+                }
+                _ => {
+                    LEADERBOARD_REQUEST_UNSUPPORTED.increment();
+                    continue;
+                }
+            },
+            ClientWorkItemKind::Reconnect => {
+                continue;
+            }
+        };
+
+        LEADERBOARD_REQUEST_OK.increment();
+
+        let stop = Instant::now();
+
+        match result {
+            Ok(_) => {
+                LEADERBOARD_RESPONSE_OK.increment();
+
+                let latency = stop.duration_since(start).as_nanos() as u64;
+
+                let _ = LEADERBOARD_RESPONSE_LATENCY.increment(latency);
+            }
+            Err(ResponseError::Exception) => {
+                LEADERBOARD_RESPONSE_EX.increment();
+            }
+            Err(ResponseError::Timeout) => {
+                LEADERBOARD_RESPONSE_TIMEOUT.increment();
+            }
+            Err(ResponseError::Ratelimited) => {
+                LEADERBOARD_RESPONSE_RATELIMITED.increment();
+            }
+            Err(ResponseError::BackendTimeout) => {
+                LEADERBOARD_RESPONSE_BACKEND_TIMEOUT.increment();
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Insert or update id-score pairs in the leaderboard.
+pub async fn upsert(
+    client: &mut LeaderboardClient,
+    config: &Config,
+    cache_name: String,
+    request: workload::leaderboard::Upsert,
+) -> std::result::Result<(), ResponseError> {
+    LEADERBOARD_UPSERT.increment();
+
+    let leaderboard = client.leaderboard(cache_name, request.leaderboard.as_ref().clone());
+    let result = timeout(
+        config.leaderboard().unwrap().request_timeout(),
+        leaderboard.upsert(request.elements),
+    )
+    .await;
+
+    record_result!(result, LEADERBOARD_UPSERT)
+}
+
+/// Get the competition rank of a list of ids in a leaderboard.
+pub async fn get_competition_rank(
+    client: &mut LeaderboardClient,
+    config: &Config,
+    cache_name: String,
+    request: workload::leaderboard::GetCompetitionRank,
+) -> std::result::Result<(), ResponseError> {
+    LEADERBOARD_GET_COMPETITION_RANK.increment();
+
+    let leaderboard = client.leaderboard(cache_name, request.leaderboard.as_ref().clone());
+    let ids = request.ids.as_ref().to_vec();
+    let result = timeout(
+        config.leaderboard().unwrap().request_timeout(),
+        leaderboard.get_competition_rank(ids),
+    )
+    .await;
+
+    record_result!(result, LEADERBOARD_GET_COMPETITION_RANK)
+}

src/clients/leaderboard/momento.rs

@@ -4,6 +4,7 @@ use momento::MomentoErrorCode;
 pub mod common;
 
 pub mod cache;
+pub mod leaderboard;
 pub mod oltp;
 pub mod ping;
 pub mod pubsub;

src/clients/mod.rs

@@ -0,0 +1,32 @@
+use super::*;
+
+#[derive(Clone, Deserialize)]
+pub struct Leaderboard {
+    /// The number of connections this process will have to each endpoint.
+    poolsize: usize,
+    /// The number of concurrent sessions per connection.
+    #[serde(default)]
+    concurrency: usize,
+    /// Request timeout
+    request_timeout: u64,
+    // number of threads for client tasks
+    threads: usize,
+}
+
+impl Leaderboard {
+    pub fn threads(&self) -> usize {
+        self.threads
+    }
+
+    pub fn poolsize(&self) -> usize {
+        std::cmp::max(1, self.poolsize)
+    }
+
+    pub fn concurrency(&self) -> usize {
+        std::cmp::max(1, self.concurrency)
+    }
+
+    pub fn request_timeout(&self) -> Duration {
+        Duration::from_millis(self.request_timeout)
+    }
+}