ref(stresstest): Distribute workloads across usecases and orgs (#91)

Makes the stresstest more realistic by writing multiple usecases and
scopes. The name of each concurrent workload is used as usecase, which
also mimics production behavior more closely. Also, each workload can
specify a number of organizations (defaulting to 1), which is used as
scope.

Internally, we create a project-level scope, since this is what the real
usecases will also be doing. We reuse the org ID as project ID.

Author: jan-auer

stresstest/example.yaml

@@ -5,6 +5,7 @@ duration: 5s
 workloads:
   - name: attachments
     mode: throughput
+    organizations: 1000
     file_sizes:
       p50: 50 KiB
       p99: 200 KiB

stresstest/src/config.rs

@@ -19,13 +19,19 @@ pub struct Workload {
     pub name: String,
     #[serde(default)]
     pub concurrency: usize,
+    #[serde(default = "default_organizations")]
+    pub organizations: u64,
     #[serde(default)]
     pub mode: WorkloadMode,
     pub file_sizes: FileSizes,
     #[serde(default)]
     pub actions: Actions,
 }
 
+fn default_organizations() -> u64 {
+    1
+}
+
 #[derive(Debug, Deserialize)]
 pub struct FileSizes {
     pub p50: ByteSize,

stresstest/src/http.rs

@@ -1,5 +1,7 @@
 //! Contains a remote implementation using HTTP to interact with objectstore.
 
+use std::collections::BTreeMap;
+
 use futures::StreamExt;
 use objectstore_client::{Client, ClientBuilder, GetResult};
 use tokio::io::AsyncReadExt;
@@ -10,23 +12,29 @@ use crate::workload::Payload;
 /// A remote implementation using HTTP to interact with objectstore.
 #[derive(Debug)]
 pub struct HttpRemote {
-    /// The Storage Client used to talk to our service.
-    pub client: Client,
+    remote: String,
+    builders: BTreeMap<String, ClientBuilder>,
 }
 
 impl HttpRemote {
     /// Creates a new `HttpRemote` instance with the given remote URL and a default client.
     pub fn new(remote: &str) -> Self {
-        let client = ClientBuilder::new(remote, "stresstest")
-            .unwrap()
-            .for_organization(12345);
-        Self { client }
+        Self {
+            remote: remote.to_owned(),
+            builders: BTreeMap::new(),
+        }
     }
 
-    pub(crate) async fn write(&self, payload: Payload) -> String {
+    pub(crate) async fn write(
+        &self,
+        usecase: &str,
+        organization_id: u64,
+        payload: Payload,
+    ) -> String {
+        let client = self.client(usecase, organization_id);
         let stream = ReaderStream::new(payload).boxed();
 
-        self.client
+        client
             .put_stream(stream)
             .compression(None)
             .send()
@@ -35,8 +43,15 @@ impl HttpRemote {
             .key
     }
 
-    pub(crate) async fn read(&self, key: &str, mut payload: Payload) {
-        let GetResult { stream, .. } = self.client.get(key).send().await.unwrap().unwrap();
+    pub(crate) async fn read(
+        &self,
+        usecase: &str,
+        organization_id: u64,
+        key: &str,
+        mut payload: Payload,
+    ) {
+        let client = self.client(usecase, organization_id);
+        let GetResult { stream, .. } = client.get(key).send().await.unwrap().unwrap();
         let mut reader = StreamReader::new(stream);
 
         // TODO: both of these are currently buffering in-memory. we should use streaming here as well.
@@ -50,7 +65,20 @@ impl HttpRemote {
         }
     }
 
-    pub(crate) async fn delete(&self, key: String) {
-        self.client.delete(&key).await.unwrap();
+    pub(crate) async fn delete(&self, usecase: &str, organization_id: u64, key: &str) {
+        let client = self.client(usecase, organization_id);
+        client.delete(key).await.unwrap();
+    }
+
+    /// Registers a new usecase that can be used by the workloads.
+    pub fn register_usecase(&mut self, usecase: &str) {
+        let builder = ClientBuilder::new(&self.remote, usecase).unwrap();
+        self.builders.insert(usecase.to_owned(), builder);
+    }
+
+    fn client(&self, usecase: &str, organization_id: u64) -> Client {
+        // NB: Reuse the organization ID as project ID to create unique projects. Right now, we do
+        // not benefit from simulating multiple projects per org.
+        self.builders[usecase].for_project(organization_id, organization_id)
     }
 }

stresstest/src/main.rs

@@ -50,6 +50,7 @@ async fn main() -> anyhow::Result<()> {
         .map(|w| {
             Workload::builder(w.name)
                 .concurrency(w.concurrency)
+                .organizations(w.organizations)
                 .mode(w.mode)
                 .size_distribution(w.file_sizes.p50.0, w.file_sizes.p99.0)
                 .action_weights(w.actions.writes, w.actions.reads, w.actions.deletes)

stresstest/src/stresstest.rs

@@ -19,8 +19,17 @@ use crate::workload::{Action, Workload, WorkloadMode};
 ///
 /// The function runs all workloads concurrently, then prints metrics and finally deletes all
 /// objects from the remote.
-pub async fn run(remote: HttpRemote, workloads: Vec<Workload>, duration: Duration) -> Result<()> {
+pub async fn run(
+    mut remote: HttpRemote,
+    workloads: Vec<Workload>,
+    duration: Duration,
+) -> Result<()> {
+    for workload in &workloads {
+        remote.register_usecase(&workload.name);
+    }
+
     let remote = Arc::new(remote);
+
     // run the workloads concurrently
     let tasks: Vec<_> = workloads
         .into_iter()
@@ -76,12 +85,12 @@ pub async fn run(remote: HttpRemote, workloads: Vec<Workload>, duration: Duratio
     let start = Instant::now();
     let cleanup_timing = Arc::new(Mutex::new(DDSketch::default()));
     futures::stream::iter(files_to_cleanup)
-        .for_each_concurrent(max_concurrency, |external_id| {
+        .for_each_concurrent(max_concurrency, |(usecase, organization_id, object_key)| {
             let remote = remote.clone();
             let cleanup_timing = cleanup_timing.clone();
             async move {
                 let start = Instant::now();
-                remote.delete(external_id).await;
+                remote.delete(&usecase, organization_id, &object_key).await;
                 cleanup_timing
                     .lock()
                     .unwrap()
@@ -147,12 +156,16 @@ async fn run_workload(
                     tokio::time::sleep(Duration::from_millis(10)).await;
                 };
 
+
                 let task = async move {
                     let start = Instant::now();
                     match action {
                         Action::Write(internal_id, payload) => {
                             let file_size = payload.len;
-                            let external_id = remote.write(payload).await;
+                            let usecase = workload.lock().unwrap().name.clone();
+                            let organization_id = workload.lock().unwrap().next_organization_id();
+                            let object_key = remote.write(&usecase, organization_id, payload).await;
+                            let external_id = (usecase, organization_id, object_key);
                             workload.lock().unwrap().push_file(internal_id, external_id);
                             let mut metrics = metrics.lock().unwrap();
                             metrics.write_timing.add(start.elapsed().as_secs_f64());
@@ -161,14 +174,16 @@ async fn run_workload(
                         }
                         Action::Read(internal_id, external_id, payload) => {
                             let file_size = payload.len;
-                            remote.read(&external_id, payload).await;
+                            let (usecase, organization_id, object_key) = &external_id;
+                            remote.read(usecase, *organization_id, object_key, payload).await;
                             workload.lock().unwrap().push_file(internal_id, external_id);
                             let mut metrics = metrics.lock().unwrap();
                             metrics.read_timing.add(start.elapsed().as_secs_f64());
                             metrics.bytes_read += file_size;
                         }
                         Action::Delete(external_id) => {
-                            remote.delete(external_id).await;
+                            let (usecase, organization_id, object_key) = &external_id;
+                            remote.delete(usecase, *organization_id, object_key).await;
                             let mut metrics = metrics.lock().unwrap();
                             metrics.delete_timing.add(start.elapsed().as_secs_f64());
                         }

stresstest/src/workload.rs

@@ -33,6 +33,7 @@ pub enum WorkloadMode {
 pub struct WorkloadBuilder {
     name: String,
     concurrency: usize,
+    organizations: u64,
     mode: WorkloadMode,
     seed: u64,
 
@@ -51,6 +52,12 @@ impl WorkloadBuilder {
         self
     }
 
+    /// The number of organizations to distribute the workload across.
+    pub fn organizations(mut self, organizations: u64) -> Self {
+        self.organizations = organizations;
+        self
+    }
+
     /// The mode of the workload, either `weighted` or `throughput`.
     pub fn mode(mut self, mode: WorkloadMode) -> Self {
         self.mode = mode;
@@ -89,6 +96,7 @@ impl WorkloadBuilder {
         Workload {
             name: self.name,
             concurrency: self.concurrency,
+            organizations: self.organizations,
             mode: self.mode,
 
             rng,
@@ -117,6 +125,8 @@ pub struct Workload {
     pub(crate) name: String,
     /// The maximum number of concurrent operations that can be performed within this workload.
     pub(crate) concurrency: usize,
+    /// The number of organizations to distribute the workload across.
+    pub(crate) organizations: u64,
     /// The target throughput for the workload, in bytes per second. Overrides concurrency.
     pub(crate) mode: WorkloadMode,
 
@@ -140,6 +150,7 @@ impl Workload {
         WorkloadBuilder {
             name: name.into(),
             concurrency: available_parallelism().unwrap().get(),
+            organizations: 1,
             mode: WorkloadMode::default(),
             seed: rand::random(),
 
@@ -233,6 +244,10 @@ impl Workload {
         }
     }
 
+    pub(crate) fn next_organization_id(&mut self) -> u64 {
+        self.rng.next_u64() % self.organizations
+    }
+
     /// Adds a file to the internal store, so it can be yielded for reads or deletes.
     ///
     /// This function has to be called for files when a write or read has completed.
@@ -262,8 +277,9 @@ impl fmt::Display for InternalId {
 
 /// Unique identifier for an object in the remote storage.
 ///
-/// These identifiers map to [`InternalId`]s.
-pub type ExternalId = String;
+/// The identifier consists of the usecase, organization ID and the object key. These key maps to
+/// [`InternalId`].
+pub type ExternalId = (String, u64, String);
 
 /// An action that can be performed by the workload.
 #[derive(Debug)]