feat: Add more stats (#4)

Author: jan-auer

src/data.rs

@@ -14,13 +14,21 @@ use time::OffsetDateTime;
 use crate::cli::{Config, MAX_PROJECTS};
 use crate::types::{SpanId, TraceId};
 
+#[derive(Clone, Debug, Default)]
+pub struct Stats {
+    pub spans: usize,
+    pub segments: usize,
+    pub traces: usize,
+}
+
 pub struct RandomGenerator<'a> {
     config: &'a Config,
     rng: ThreadRng,
     segment_dist: Normal<f64>,
     span_dist: Normal<f64>,
     #[allow(dead_code, reason = "TODO: support custom receive time")]
     batch_delay_dist: Normal<f64>,
+    stats: Stats,
 }
 
 impl<'a> RandomGenerator<'a> {
@@ -49,9 +57,14 @@ impl<'a> RandomGenerator<'a> {
             segment_dist,
             span_dist,
             batch_delay_dist,
+            stats: Stats::default(),
         }
     }
 
+    pub fn stats(&self) -> &Stats {
+        &self.stats
+    }
+
     pub fn rng(&mut self) -> &mut ThreadRng {
         &mut self.rng
     }
@@ -73,10 +86,12 @@ impl<'a> RandomGenerator<'a> {
     }
 
     pub fn trace(&mut self) -> TraceInfo {
+        self.stats.traces += 1;
         TraceInfo::new(self.organization_id())
     }
 
     pub fn segment<'b>(&mut self, trace: &'b TraceInfo) -> SegmentInfo<'b> {
+        self.stats.segments += 1;
         SegmentInfo::new(trace, self.project_id(trace.organization_id))
     }
 
@@ -115,6 +130,8 @@ impl<'a> RandomGenerator<'a> {
     }
 
     pub fn span<'b>(&mut self, segment: &'b SegmentInfo<'b>, span_ref: SpanRef) -> Span {
+        self.stats.spans += 1;
+
         let now = OffsetDateTime::now_utc();
         let lower = now - Duration::from_secs(60 * 60);
 

src/main.rs

@@ -15,36 +15,28 @@ mod types;
 
 struct StdoutProducer {
     stdout: StdoutLock<'static>,
-    count: usize,
 }
 
 impl StdoutProducer {
     pub fn new() -> Self {
         Self {
             stdout: std::io::stdout().lock(),
-            count: 0,
         }
     }
 
     pub fn produce_json<T: Serialize>(&mut self, value: &T) -> Result<()> {
         serde_json::to_writer(&mut self.stdout, value)?;
         writeln!(&mut self.stdout)?;
-
-        self.count += 1;
         Ok(())
     }
-
-    pub fn count(&self) -> usize {
-        self.count
-    }
 }
 
 fn produce(config: &Config) -> Result<()> {
     let start = Instant::now();
     let mut generator = RandomGenerator::new(config);
     let mut producer = StdoutProducer::new();
 
-    while producer.count() < config.count {
+    while generator.stats().spans < config.count {
         let trace = generator.trace();
         let mut remote_parent = None;
 
@@ -69,11 +61,11 @@ fn produce(config: &Config) -> Result<()> {
         }
     }
 
-    log::info!(
-        "Produced {} spans in {:?}",
-        producer.count(),
-        start.elapsed()
-    );
+    let stats = generator.stats();
+    log::info!("Finished in {:?}", start.elapsed());
+    log::info!("  traces:   {}", stats.traces);
+    log::info!("  segments: {}", stats.segments);
+    log::info!("  spans:    {}", stats.spans);
 
     Ok(())
 }