feat: Support throughput throttling (#5)

Adds an option to throttle the throughput of spans per second. This
option only makes sense for larger volumes. We will not throttle if the
requested throughput is higher than what spangen can produce, but
there's also no warning.

Author: jan-auer

.github/workflows/ci.yml

@@ -20,7 +20,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install Rust Toolchain
-        run: rustup toolchain install stable --profile minimal --component rustfmt clippy --no-self-update
+        run: rustup toolchain install stable --profile minimal --component rustfmt,clippy --no-self-update
 
       - name: Run rustfmt
         run: cargo fmt --check

README.md

@@ -6,7 +6,13 @@ A load generator for spans in traces that adheres to the `snuba-spans` schema.
 
 ```
 Options:
-  --count           the number of spans to generate in total.
+  --count           the number of spans to generate in total. spangen will stop
+                    generating new traces after this number has been reached,
+                    but it will finish started traces and segments. The actual
+                    number of spans generated may therefore be higher than this
+                    option.
+  --throughput      the throughput of spans per second (defaults to no
+                    throttling).
   --spans-per-segment
                     the average number of spans per segment (randomized).
   --spans-per-segment-stddev
@@ -37,9 +43,8 @@ Options:
   --segments-without-root
                     the percentage of segments without an explicit root span
                     (0..100)
-  --number-of-orgs  the number of organizations.
-  --number-of-projects
-                    the number of projects per organization.
+  --orgs            the number of organizations.
+  --projects        the number of projects per organization.
   --help, help      display usage information
 ```
 

src/cli.rs

@@ -7,9 +7,17 @@ pub const MAX_PROJECTS: u64 = 1000;
 #[derive(Debug, FromArgs)]
 pub struct Config {
     /// the number of spans to generate in total.
+    ///
+    /// spangen will stop generating new traces after this number has been reached, but it will
+    /// finish started traces and segments. The actual number of spans generated may therefore be
+    /// higher than this option.
     #[argh(option)]
     pub count: usize,
 
+    /// the throughput of spans per second (defaults to no throttling).
+    #[argh(option)]
+    pub throughput: Option<u32>,
+
     /// the average number of spans per segment (randomized).
     #[argh(option, default = "17")]
     pub spans_per_segment: usize,

src/main.rs

@@ -1,5 +1,6 @@
 use std::io::{StdoutLock, Write};
-use std::time::Instant;
+use std::thread;
+use std::time::{Duration, Instant};
 
 use anyhow::Result;
 use rand::Rng;
@@ -13,6 +14,47 @@ mod cli;
 mod data;
 mod types;
 
+const MIN_SLEEP: Duration = Duration::from_millis(1);
+
+struct Throttle {
+    throughput: Option<u32>,
+    accepted: u32,
+    last_sleep: Instant,
+}
+
+impl Throttle {
+    pub fn new(throughput: Option<u32>) -> Self {
+        Self {
+            throughput,
+            accepted: 0,
+            last_sleep: Instant::now(),
+        }
+    }
+
+    pub fn accept(&mut self) {
+        if self.throughput.is_some() {
+            self.accepted += 1;
+        }
+    }
+
+    pub fn wait(&mut self) {
+        let Some(throughput) = self.throughput else {
+            return;
+        };
+
+        let now = Instant::now();
+        let elapsed = now - self.last_sleep;
+
+        let expected_duration = (Duration::from_secs(1) * self.accepted) / throughput;
+        let sleep_duration = expected_duration.saturating_sub(elapsed);
+        if sleep_duration >= MIN_SLEEP {
+            thread::sleep(sleep_duration);
+            self.last_sleep = now + sleep_duration;
+            self.accepted = 0;
+        }
+    }
+}
+
 struct StdoutProducer {
     stdout: StdoutLock<'static>,
 }
@@ -35,6 +77,7 @@ fn produce(config: &Config) -> Result<()> {
     let start = Instant::now();
     let mut generator = RandomGenerator::new(config);
     let mut producer = StdoutProducer::new();
+    let mut throttle = Throttle::new(config.throughput);
 
     while generator.stats().spans < config.count {
         let trace = generator.trace();
@@ -44,6 +87,8 @@ fn produce(config: &Config) -> Result<()> {
             let segment = generator.segment(&trace);
             let span_refs = generator.span_refs();
 
+            throttle.wait();
+
             for span_ref in &span_refs {
                 let mut span = generator.span(&segment, *span_ref);
                 if span_ref.parent_id.is_none() {
@@ -52,6 +97,7 @@ fn produce(config: &Config) -> Result<()> {
                 }
 
                 producer.produce_json(&span)?;
+                throttle.accept();
             }
 
             // in 50% of the cases, pick a random span as the remote parent for the next segment