chore(antithesis): Bound contexts in runs

This commit introduces a new protocol to the intake node to allow for
context bounding in runs. The purposes are twofold. First, while it's
interesting to explore timelines in which there are too many contexts
for either lane to handle these are less relevant to me _now_, so I've
capped total contexts at 1M. Second, in order for timeseries equality
to _work_ I need to be certain that I have at least _some_ timeseries
with more than one point emitted into them. Previously drivers were
emitting essentially random contexts and it was very unlikely for more
than one point to be emitted per context. Oops.

Drivers now request N contexts from the intake service for each sub-kind
of dogstatsd they emit, meaning all contexts are sourced from a single
spot in a topology and not from pure randomness.

Author: blt

Cargo.lock

@@ -20,6 +20,7 @@ itoa = { workspace = true }
 libc = { workspace = true }
 num-traits = { workspace = true }
 rand = { workspace = true }
+reqwest = { workspace = true, features = ["blocking"] }
 ryu = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }

test/antithesis/harness/Cargo.toml

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 671f360ca24ca5add047c2f93acf185116ebfa9c663637a430a15fa6abc56151 # shrinks to seed = 8436120490560261333, kind = Metric

test/antithesis/harness/proptest-regressions/contexts.txt

@@ -14,7 +14,7 @@ use antithesis_sdk::prelude::*;
 use antithesis_sdk::random::AntithesisRng;
 use anyhow::Context;
 use clap::Parser;
-use harness::config::DatadogConfig;
+use harness::config::{ContextSourceConfig, DatadogConfig};
 use rand::rand_core::UnwrapErr;
 use serde_json::json;
 
@@ -63,10 +63,22 @@ fn main() -> anyhow::Result<()> {
 
     // Load-generator view of the same sample, so a generator caps its datagrams
     // to the SUT's receive buffer without reading the Agent config.
+    let driver = config.driver_config(&mut rng);
     let driver_path = cli.config_dir.join("driver.yaml");
-    fs::write(&driver_path, config.driver_config(&mut rng).to_yaml()?.as_bytes())
+    fs::write(&driver_path, driver.to_yaml()?.as_bytes())
         .with_context(|| format!("write driver config {}", driver_path.display()))?;
 
+    // Per-kind context-pool caps for this timeline, read lazily by the intake's shared pool. Sampled
+    // here so the pool's cardinality varies per timeline like every other sampled knob.
+    let context_source_path = cli.config_dir.join("context_source.yaml");
+    fs::write(
+        &context_source_path,
+        ContextSourceConfig::sample(&mut rng, driver.payload_byte_limit)
+            .to_yaml()?
+            .as_bytes(),
+    )
+    .with_context(|| format!("write context source config {}", context_source_path.display()))?;
+
     // Per-timeline anchor: counting these in triage tells us how many distinct
     // configs the run sampled.
     let details = serde_json::to_value(&config).unwrap_or_else(|e| json!({ "serialize_error": e.to_string() }));

test/antithesis/harness/src/bin/first_sample_config/main.rs

@@ -245,6 +245,13 @@ cloud_provider_metadata: []
 /// Upper bound on datagrams one driver invocation ships in a timeline.
 const MAX_DATAGRAMS: usize = 10_000;
 
+/// Upper bound on the working set one driver invocation fetches from the shared context pool.
+const MAX_WORKING_SET: u64 = 1_024;
+
+/// The intake's ceiling on one `/contexts` request. A `context_count` past it is rejected there, so a
+/// config carrying one is rejected here instead.
+const MAX_CONTEXTS_PER_REQUEST: usize = 65_536;
+
 /// Config a load generator reads to shape its output to this timeline's SUT.
 /// `first_sample_config` samples it beside `datadog.yaml` from one draw, so the
 /// generator and the SUT are driven together.
@@ -256,6 +263,14 @@ pub struct DriverConfig {
     pub payload_byte_limit: usize,
     /// Datagrams a driver invocation ships this timeline.
     pub datagram_count: usize,
+    /// Distinct contexts a driver invocation fetches from the shared pool as its working set.
+    ///
+    /// Sampled boundary-biased log-uniform in `1..=1_024`. Valid values run `1..=65_536`, the intake's
+    /// per-request ceiling. Outside that the intake rejects every `/contexts` request, the driver waits
+    /// out its fetch budget and ships nothing, so [`Self::read`] rejects such a config rather than
+    /// letting a timeline generate no load. A larger working set spreads load over more identities and
+    /// so puts fewer points in each, which is the trade against recurrence.
+    pub context_count: usize,
 }
 
 impl DriverConfig {
@@ -270,6 +285,7 @@ impl DriverConfig {
         Self {
             payload_byte_limit,
             datagram_count: rng.random_range(0..=MAX_DATAGRAMS),
+            context_count: usize::try_from(Probe::new(1, MAX_WORKING_SET).sample(rng)).unwrap_or(usize::MAX),
         }
     }
 
@@ -292,10 +308,96 @@ impl DriverConfig {
     pub fn read(config_dir: &Path) -> anyhow::Result<Self> {
         let path = config_dir.join("driver.yaml");
         let yaml = fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?;
-        serde_yaml::from_str(&yaml).with_context(|| format!("parse driver config from {}", path.display()))
+        let config: Self =
+            serde_yaml::from_str(&yaml).with_context(|| format!("parse driver config from {}", path.display()))?;
+        anyhow::ensure!(
+            (1..=MAX_CONTEXTS_PER_REQUEST).contains(&config.context_count),
+            "context_count {} outside 1..={} in {}, the intake would reject every fetch and the driver would ship nothing",
+            config.context_count,
+            MAX_CONTEXTS_PER_REQUEST,
+            path.display()
+        );
+        Ok(config)
+    }
+}
+
+/// Upper bound on the distinct contexts a shared pool holds across every kind. The pool retains each
+/// minted context, so the ceiling belongs to the total rather than to any one kind.
+const MAX_CONTEXTS_TOTAL: u64 = 1_000_000;
+
+/// The per-kind caps a timeline's shared context pool fills to before it recurs existing contexts.
+/// `first_sample_config` samples this beside `datadog.yaml` so cardinality varies per timeline. Each
+/// cap is drawn against the budget the earlier draws left, so a kind's cardinality still varies at
+/// random while the three together stay under [`MAX_CONTEXTS_TOTAL`].
+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+pub struct ContextSourceConfig {
+    /// Bytes a rendered line of any pooled context must fit, this timeline's real datagram budget
+    /// rather than the protocol ceiling. Mint builds identities against it, so every served context
+    /// has a rendering the driver can pack. Defaults to the sampled `payload_byte_limit`, which is the
+    /// smaller of the SUT's receive buffer and [`PAYLOAD_BYTE_LIMIT`].
+    pub payload_byte_limit: usize,
+    /// Distinct metric contexts the pool holds before it recurs the ones it has.
+    ///
+    /// Sampled in `1..=MAX_CONTEXTS_TOTAL` minus what the other kinds took. A larger cap explores more
+    /// identities and puts fewer points in each, and costs memory: the pool retains every context it
+    /// mints for the life of the run. Zero is not sampled and serves nothing for the kind.
+    pub metric_contexts: usize,
+    /// Distinct event contexts the pool holds. Same range and trade as [`Self::metric_contexts`].
+    pub event_contexts: usize,
+    /// Distinct service-check contexts the pool holds. Same range and trade as
+    /// [`Self::metric_contexts`].
+    pub service_check_contexts: usize,
+}
+
+impl ContextSourceConfig {
+    /// Sample the per-kind caps, each boundary-biased log-uniform against the budget still free.
+    ///
+    /// The draws run in order and each spends from one shared ceiling, so the total is bounded by
+    /// construction rather than by scaling three independent draws afterwards. Every kind keeps at
+    /// least one context, since a kind capped at zero panics the pool's draw.
+    #[must_use]
+    pub fn sample<R: Rng + ?Sized>(rng: &mut R, payload_byte_limit: usize) -> Self {
+        // Two contexts held back so the later kinds can each keep their one.
+        let metric_contexts = sample_cap(rng, MAX_CONTEXTS_TOTAL - 2);
+        let free = MAX_CONTEXTS_TOTAL - metric_contexts as u64;
+        let event_contexts = sample_cap(rng, free - 1);
+        let free = free - event_contexts as u64;
+        Self {
+            payload_byte_limit,
+            metric_contexts,
+            event_contexts,
+            service_check_contexts: sample_cap(rng, free),
+        }
+    }
+
+    /// Render `self` as a `context_source.yaml` string.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if serialization fails.
+    pub fn to_yaml(&self) -> anyhow::Result<String> {
+        serde_yaml::to_string(self).context("serialize context_source.yaml")
+    }
+
+    /// Read the context-source config from the `context_source.yaml` that `first_sample_config` wrote
+    /// to `config_dir`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the config is unreadable or is not valid YAML.
+    pub fn read(config_dir: &Path) -> anyhow::Result<Self> {
+        let path = config_dir.join("context_source.yaml");
+        let yaml = fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?;
+        serde_yaml::from_str(&yaml).with_context(|| format!("parse context source config from {}", path.display()))
     }
 }
 
+/// A single per-kind cap in `1..=ceiling`. The ceiling is well within `usize` on every supported
+/// target, so the saturating conversion is unreachable in practice.
+fn sample_cap<R: Rng + ?Sized>(rng: &mut R, ceiling: u64) -> usize {
+    usize::try_from(Probe::new(1, ceiling.max(1)).sample(rng)).unwrap_or(usize::MAX)
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeSet;
@@ -401,6 +503,28 @@ mod tests {
         assert_eq!(seen, [true, true]);
     }
 
+    // The pool holds every minted context, so the ceiling is on the total across kinds rather than on
+    // each kind alone. Three independent draws at the ceiling would retain three million.
+    #[test]
+    fn context_caps_sum_within_the_total_ceiling() {
+        for seed in 0..64 {
+            let caps = ContextSourceConfig::sample(&mut SeqRng(seed), 8_192);
+            let total = (caps.metric_contexts + caps.event_contexts + caps.service_check_contexts) as u64;
+            assert!(total <= MAX_CONTEXTS_TOTAL, "seed {seed} sampled {total}");
+            // A kind of zero panics the pool's `random_range(0..0)`, so every kind keeps at least one.
+            assert!(caps.metric_contexts >= 1 && caps.event_contexts >= 1 && caps.service_check_contexts >= 1);
+        }
+    }
+
+    // Randomness still drives each kind rather than the total being split evenly.
+    #[test]
+    fn context_caps_vary_per_kind() {
+        let spread: BTreeSet<usize> = (0..64)
+            .map(|seed| ContextSourceConfig::sample(&mut SeqRng(seed), 8_192).metric_contexts)
+            .collect();
+        assert!(spread.len() > 8, "metric cap barely varies: {spread:?}");
+    }
+
     #[test]
     fn compressor_samples_every_kind() {
         let mut seen = BTreeSet::new();