diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 620e7f90..20e5b103 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,6 +22,13 @@ jobs: runs-on: ubuntu-latest steps: + # librdkafka is built from source and links against libcurl, whose headers are not + # on the runner by default. Needed by anything that builds the `kafka` feature. + - name: Install libcurl-dev + run: | + sudo apt-get update + sudo apt-get install -y libcurl4-openssl-dev + - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 - name: Install Rust Toolchain @@ -75,6 +82,13 @@ jobs: runs-on: ubuntu-latest steps: + # librdkafka is built from source and links against libcurl, whose headers are not + # on the runner by default. Needed by anything that builds the `kafka` feature. + - name: Install libcurl-dev + run: | + sudo apt-get update + sudo apt-get install -y libcurl4-openssl-dev + - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 - name: Start Bigtable Emulator @@ -181,6 +195,13 @@ jobs: id-token: write steps: + # librdkafka is built from source and links against libcurl, whose headers are not + # on the runner by default. Needed by anything that builds the `kafka` feature. + - name: Install libcurl-dev + run: | + sudo apt-get update + sudo apt-get install -y libcurl4-openssl-dev + - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 - name: Install Rust Toolchain diff --git a/Cargo.lock b/Cargo.lock index 2224f970..96a4d630 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -249,6 +249,18 @@ dependencies = [ "serde", ] +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-compression" version = "0.4.42" @@ -491,6 +503,20 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -679,6 +705,12 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "convert_case" version = "0.10.0" @@ -1099,7 +1131,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1778,7 +1810,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "system-configuration", "tokio", "tower-service", @@ -2197,6 +2229,18 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libz-sys" +version = "1.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2459,7 +2503,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2574,6 +2618,28 @@ dependencies = [ "libc", ] +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "objc2" version = "0.6.4" @@ -2768,6 +2834,18 @@ dependencies = [ "zstd-safe", ] +[[package]] +name = "objectstore-inventory-tracker" +version = "0.1.0" +dependencies = [ + "blake3", + "rdkafka", + "serde", + "serde_json", + "thiserror", + "tracing", +] + [[package]] name = "objectstore-log" version = "0.1.0" @@ -3276,6 +3354,15 @@ dependencies = [ "elliptic-curve", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -3455,7 +3542,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.6.4", + "socket2 0.5.10", "thiserror", "tokio", "tracing", @@ -3493,7 +3580,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.4", + "socket2 0.5.10", "tracing", "windows-sys 0.60.2", ] @@ -3642,6 +3729,39 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rdkafka" +version = "0.39.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7956f9ac12b5712e50372d9749a3102f4810a8d42481c5eae3748d36d585bcf" +dependencies = [ + "futures-channel", + "futures-util", + "libc", + "log", + "rdkafka-sys", + "serde", + "serde_derive", + "serde_json", + "slab", + "tokio", + "tracing", +] + +[[package]] +name = "rdkafka-sys" +version = "4.10.0+2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e234cf318915c1059d4921ef7f75616b5219b10b46e9f3a511a15eb4b56a3f77" +dependencies = [ + "cmake", + "libc", + "libz-sys", + "num_enum", + "openssl-sys", + "pkg-config", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3886,7 +4006,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3945,7 +4065,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4080,7 +4200,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b55fb86dfd3a2f5f76ea78310a88f96c4ea21a3031f8d212443d56123fd0521" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4514,7 +4634,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -4645,7 +4765,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4836,6 +4956,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap 2.14.0", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + [[package]] name = "tonic" version = "0.14.6" @@ -5432,7 +5582,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5667,6 +5817,15 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/Cargo.toml b/Cargo.toml index 3ae8b886..820365e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ debug = true lto = "thin" [workspace.dependencies] +objectstore-inventory-tracker = { path = "objectstore-inventory-tracker" } objectstore-client = { path = "clients/rust" } objectstore-log = { path = "objectstore-log" } objectstore-metrics = { path = "objectstore-metrics" } @@ -42,6 +43,7 @@ async-trait = "0.1.89" axum = "0.8.9" base64 = "0.22.1" bigtable_rs = { git = "https://github.com/getsentry/bigtable_rs.git", rev = "4cb75bc5e5f87204363973f6302107768e64972e" } +blake3 = "1.8.5" bytes = "1.12.0" bytesize = "2.4.0" chrono = "0.4.45" @@ -81,6 +83,10 @@ rustls = { version = "0.23.40", default-features = false } secrecy = "0.10.3" sentry = "0.48.3" sentry-options = "1.2.4" +# - `ssl` is required for SASL/SCRAM +# - `librdkafka` must be built from source either way, and `cmake-build` is the +# recommended way to do it +rdkafka = { version = "0.39.0", features = ["cmake-build", "ssl", "tracing"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150" serde_yaml = "0.9.34-deprecated" diff --git a/objectstore-inventory-tracker/Cargo.toml b/objectstore-inventory-tracker/Cargo.toml new file mode 100644 index 00000000..a2863951 --- /dev/null +++ b/objectstore-inventory-tracker/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "objectstore-inventory-tracker" +authors = ["Sentry "] +description = "Emits per-record inventory change events for shared storage resources" +homepage = "https://getsentry.github.io/objectstore/" +repository = "https://github.com/getsentry/objectstore" +license-file = "../LICENSE.md" +version = "0.1.0" +edition = "2024" +publish = false + +[features] +# Off by default so that depending on the record types does not pull in librdkafka +# and its native build. +kafka = ["dep:rdkafka"] +# Exposes `DummyProducer` so downstream crates can assert on emitted records +# without a broker. +test-utils = [] + +[dependencies] +blake3 = { workspace = true } +# Allow a range of possible `rdkafka` versions so library users can pin whichever they use +rdkafka = { version = ">=0.29.0, <0.40.0", features = [ + "cmake-build", + "ssl", + "tracing", +], optional = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +serde_json = { workspace = true } diff --git a/objectstore-inventory-tracker/README.md b/objectstore-inventory-tracker/README.md new file mode 100644 index 00000000..159ac8ca --- /dev/null +++ b/objectstore-inventory-tracker/README.md @@ -0,0 +1,58 @@ +# objectstore-inventory-tracker + +Wrapper around a Kafka producer that a storage service can use to emit a +lightweight change stream. + +Produces to Kafka via [`sentry_arroyo`](https://docs.rs/sentry_arroyo/). +Message schema is defined in [sentry-kafka-schemas](https://github.com/getsentry/sentry-kafka-schemas). + +## Why + +Generalized cheap analytics. A stream consumer can merge each change message +into an external table to maintain an inventory of each record stored along with +its size and some attribution metadata. Various analyses (usage trends, Cost of +Goods Sold (COGS) rollups) can be run on the external table relatively cheaply +and without putting stress on the datastores actually used to serve production +traffic. + +`InventoryTracker` is not meant to be used as a write-ahead log or subscription +channel for record changes. It was designed with back-of-house analytics in +mind; record IDs are hashed to enable sampling and strip any possible sensitive +information, and stored data is not included in the stream. + +## Usage + +```rust +use objectstore_inventory_tracker::{InventoryTracker, KafkaConfig, KafkaProducer}; +use std::time::SystemTime; + +// Create your Kafka producer. +let producer = KafkaProducer::try_new(KafkaConfig { + topic: "shared-resources-inventory".into(), + bootstrap_servers: vec!["localhost:9092".into()], + override_params: Default::default(), +})?; + +// Create your `InventoryTracker`. This one is for the `my_gcs_bucket` bucket +// and has a sampling rate of `1.0` (unsampled). +let tracker = InventoryTracker::new(producer, "my_gcs_bucket", 1.0); + +// Emit a message indicating that the given key has been written with a size of +// 4096 bytes. Hashing and sampling of the key are handled internally; if the key +// is not sampled, this is a no-op. +let storage_key = "attachments/org.1/project.1/objects/abc"; +tracker.write(storage_key, "attachments", 4096, SystemTime::now(), None, Some(1), Some(1))?; +``` + +The `"my_gcs_bucket"` string in the above example is meant to correspond to a +label on the specific GCS bucket your service uses to store data. That way, an +inventory derived from your change stream can be joined with billing data to +analyze costs. If you have multiple buckets, or multiple storage backends, it's +recommended that you configure each of them with their own `InventoryTracker`. + +### Features + +| Feature | Description | +|---|---| +| `kafka` | The `sentry_arroyo`-backed producer. Off by default, so the record types are usable without pulling in arroyo, librdkafka, and their native build. | +| `test-utils` | Exports `DummyProducer`, so downstream crates can assert on emitted records without a broker. | diff --git a/objectstore-inventory-tracker/src/kafka.rs b/objectstore-inventory-tracker/src/kafka.rs new file mode 100644 index 00000000..a61b6f86 --- /dev/null +++ b/objectstore-inventory-tracker/src/kafka.rs @@ -0,0 +1,230 @@ +//! A [`Producer`] backed by [`rdkafka`]. + +use std::collections::HashMap; +use std::time::Duration; + +use rdkafka::ClientConfig; +use rdkafka::client::ClientContext; +use rdkafka::error::KafkaError; +use rdkafka::producer::{ + BaseRecord, DeliveryResult, Producer as RdKafkaProducer, ProducerContext, ThreadedProducer, +}; + +use crate::producer::Producer; + +/// Connection settings for [`KafkaProducer`]. +#[derive(Clone, Debug)] +pub struct KafkaConfig { + /// Topic to produce to. + pub topic: String, + + /// Broker addresses. + pub bootstrap_servers: Vec, + + /// Additional librdkafka properties, passed through verbatim. + pub override_params: HashMap, +} + +/// Something that went wrong producing to Kafka. +#[derive(Debug, thiserror::Error)] +pub enum KafkaProducerError { + /// The producer could not be created from the given configuration. + #[error("failed to create kafka producer")] + InvalidConfig(#[source] KafkaError), + + /// The record could not be enqueued. + /// + /// This may mean the local queue is full, which may be caused by the broker being + /// unreachable or backed up. + #[error("failed to enqueue inventory record")] + SendFailed(#[source] KafkaError), + + /// The local queue was not emptied before the flush timeout elapsed. + #[error("failed to flush inventory records")] + FlushFailed(#[source] KafkaError), +} + +/// Called for each record that fails to deliver. +/// +/// Boxed rather than generic so that [`KafkaProducer::try_new`] can take an `Option` of +/// it: `None` would otherwise need a turbofish, since there is no closure to infer a type +/// parameter from. +/// +/// Invoked from a librdkafka background thread, so it must not block. +pub type OnDeliveryFailure = Box; + +/// Reports delivery outcomes. +/// +/// Delivery is asynchronous, so a successful [`Producer::send`] only means the record was +/// enqueued locally. Without a delivery callback, a broker rejecting records looks exactly +/// like everything working. +struct DeliveryReporter { + on_failure: Option, +} + +impl ClientContext for DeliveryReporter {} + +impl ProducerContext for DeliveryReporter { + type DeliveryOpaque = (); + + fn delivery(&self, result: &DeliveryResult<'_>, _opaque: Self::DeliveryOpaque) { + if let Err((error, _)) = result { + tracing::warn!(%error, "failed to deliver inventory record"); + if let Some(on_failure) = &self.on_failure { + on_failure(error); + } + } + } +} + +/// Produces inventory records onto a Kafka topic. +/// +/// Sends are non-blocking: records go onto librdkafka's internal queue and a background +/// thread delivers them. When that queue is full, [`Producer::send`] returns an error +/// rather than waiting. +pub struct KafkaProducer { + topic: String, + producer: ThreadedProducer, +} + +// Hand-written because `ThreadedProducer` is not `Debug`. +impl std::fmt::Debug for KafkaProducer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KafkaProducer") + .field("topic", &self.topic) + .field("in_flight", &self.producer.in_flight_count()) + .finish_non_exhaustive() + } +} + +impl KafkaProducer { + /// Creates a producer from `config`, optionally reporting delivery failures to + /// `on_delivery_failure`. + /// + /// Failures are always logged. The callback is how a caller additionally counts them, + /// without this crate having to depend on a metrics backend; pass `None` to only log. + pub fn try_new( + config: KafkaConfig, + on_delivery_failure: Option, + ) -> Result { + let mut client_config = ClientConfig::new(); + client_config.set("bootstrap.servers", config.bootstrap_servers.join(",")); + + // Applied after the broker list so that a caller can override it if they need to. + for (key, value) in &config.override_params { + client_config.set(key, value); + } + + let producer = client_config + .create_with_context(DeliveryReporter { + on_failure: on_delivery_failure, + }) + .map_err(KafkaProducerError::InvalidConfig)?; + + Ok(Self { + topic: config.topic, + producer, + }) + } +} + +impl Producer for KafkaProducer { + type Error = KafkaProducerError; + + fn send(&self, key: &[u8], payload: Vec) -> Result<(), Self::Error> { + let record: BaseRecord<'_, [u8], [u8]> = + BaseRecord::to(&self.topic).key(key).payload(&payload); + + self.producer + .send(record) + .map_err(|(error, _)| KafkaProducerError::SendFailed(error)) + } + + fn flush(&self, timeout: Duration) -> Result<(), Self::Error> { + self.producer + .flush(timeout) + .map_err(KafkaProducerError::FlushFailed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> KafkaConfig { + KafkaConfig { + topic: "shared-resources-inventory".into(), + bootstrap_servers: vec!["127.0.0.1:9092".into()], + override_params: HashMap::from([("compression.type".into(), "lz4".into())]), + } + } + + #[test] + fn producer_targets_the_configured_topic() { + let producer = KafkaProducer::try_new(config(), None).unwrap(); + assert_eq!(producer.topic, "shared-resources-inventory"); + } + + #[test] + fn every_bootstrap_server_is_passed_through() { + let mut config = config(); + config.bootstrap_servers = vec!["a:9092".into(), "b:9092".into(), "c:9092".into()]; + + // librdkafka parses the broker list at creation, so building successfully is the + // assertion that the joined value was well formed. + assert!(KafkaProducer::try_new(config, None).is_ok()); + } + + #[test] + fn bad_config_is_an_error() { + let mut config = config(); + config + .override_params + .insert("not.a.real.property".into(), "1".into()); + + assert!(matches!( + KafkaProducer::try_new(config, None), + Err(KafkaProducerError::InvalidConfig(_)) + )); + } + + #[test] + fn flushing_an_empty_queue_returns_immediately() { + let producer = KafkaProducer::try_new(config(), None).unwrap(); + + let start = std::time::Instant::now(); + producer.flush(Duration::from_secs(5)).unwrap(); + assert!(start.elapsed() < Duration::from_secs(1)); + } + + #[test] + fn delivery_failures_reach_the_callback() { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let failures = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&failures); + + let mut config = config(); + // Nothing is listening on this port, and a short timeout means librdkafka gives + // up and reports the record as undeliverable while the test is still running. + config.bootstrap_servers = vec!["127.0.0.1:1".into()]; + config + .override_params + .insert("message.timeout.ms".into(), "300".into()); + + let producer = KafkaProducer::try_new( + config, + Some(Box::new(move |_| { + counter.fetch_add(1, Ordering::SeqCst); + })), + ) + .unwrap(); + + producer.send(b"key", b"payload".to_vec()).unwrap(); + // Flushing drives the delivery callbacks; the record cannot succeed. + let _ = producer.flush(Duration::from_secs(5)); + + assert_eq!(failures.load(Ordering::SeqCst), 1); + } +} diff --git a/objectstore-inventory-tracker/src/lib.rs b/objectstore-inventory-tracker/src/lib.rs new file mode 100644 index 00000000..130feff2 --- /dev/null +++ b/objectstore-inventory-tracker/src/lib.rs @@ -0,0 +1,91 @@ +//! Emits per-record inventory change events for shared storage resources. +//! +//! Change events are published to the `shared-resources-inventory` Kafka topic. They +//! include (among other things): +//! - `shared_resource_id`: identifies the storage backend the record is hosted in. This +//! is meant to be, for example, a specific GCS bucket or Bigtable instance. +//! - `record_id`: identifies each record. `InventoryTracker` populates this with a hash +//! of the identifier passed in by the caller. +//! - `size`: the size of the record in bytes (including metadata). +//! - `expiration_time`: a timestamp (unixtime microseconds) describing when the record is +//! meant to be deleted. +//! +//! # Usage +//! +//! ``` +//! use objectstore_inventory_tracker::{InventoryTracker, NoopProducer}; +//! use std::time::SystemTime; +//! +//! # let producer = NoopProducer; +//! // Create your `InventoryTracker`. This one is for the `my_gcs_bucket` bucket +//! // and has a sampling rate of `1.0` (unsampled). +//! let tracker = InventoryTracker::new(producer, "my_gcs_bucket", 1.0); +//! +//! // Emit a message indicating that `storage_key` has been written, 4096 bytes in size. +//! // Sampling and hashing of the key are handled internally. +//! let storage_key = "example_feature/org.1/project.1/objects/abc"; +//! tracker.write(storage_key, "example_feature", 4096, SystemTime::now(), None, Some(1), Some(1))?; +//! # Ok::<(), std::convert::Infallible>(()) +//! ``` +//! +//! The `"my_gcs_bucket"` string in the above example is meant to correspond to a +//! label on the specific GCS bucket your service uses to store data. That way, an +//! inventory derived from your change stream can be joined with billing data to +//! analyze costs. If you have multiple buckets, or multiple storage backends, it's +//! recommended that you configure each of them with their own `InventoryTracker`. +//! +//! # Sampling +//! +//! [`InventoryTracker`] hashes the storage key it gets from the caller and uses part of +//! the digest to determine whether change messages should be emitted for a given object. +//! Each [`InventoryTracker`] instance is configured with its own sample rate so services +//! can enable, disable, or tune sampling as needed. +//! +//! Each change message includes the sampling rate that was in effect when the message was +//! emitted. Downstream consumers can use `1 / sample_rate` as a weight when computing +//! aggregates to approximate what the unsampled aggregate would have been. +//! +//! When sampling is used the resulting dataset will generally be representative. However, +//! it will not be complete, and with a low sample rate it's more likely that a small +//! trend or subpopulation of your data will be entirely missing from the sampled dataset. +//! +//! # Fail open +//! +//! [`Producer::send`] does not block, and it returns an error instead of waiting when the +//! local queue is full. Callers may count the error and move on. +//! +//! Sending is asynchronous, so messages handed over just before a process exits are still +//! sitting in a local queue. Call [`InventoryTracker::flush`] during shutdown to deliver +//! them; anything still queued when the process goes away is lost. +//! +//! Losing change messages causes drift in downstream consumers. If your storage service +//! expires data, this drift will sort itself out as records age out. If your service +//! retains data indefinitely, consider writing a periodic reconciliation job. +//! +//! # Ordering +//! +//! Each change message has a `timestamp` field with microsecond resolution so that +//! multiple changes to the same record are ordered. +//! +//! NOTE: This crate can provide no guarantees that this ordering scheme will always +//! result in messages being serialized in Kafka in the same order the corresponding +//! operations were serialized in your service or its storage layer. If this is an issue, +//! consider writing a periodic reconciliation job or implementing synchronization outside +//! of this crate. + +#![warn(missing_docs)] +#![warn(missing_debug_implementations)] + +mod producer; +mod record; +mod tracker; + +#[cfg(feature = "kafka")] +pub mod kafka; + +pub use producer::{NoopProducer, Producer}; +pub use record::{InventoryRecord, OpType, epoch_micros}; +pub use tracker::InventoryTracker; + +#[cfg(any(test, feature = "test-utils"))] +pub use producer::{DummyProducer, SentMessage}; diff --git a/objectstore-inventory-tracker/src/producer.rs b/objectstore-inventory-tracker/src/producer.rs new file mode 100644 index 00000000..44e6da0d --- /dev/null +++ b/objectstore-inventory-tracker/src/producer.rs @@ -0,0 +1,122 @@ +//! The transport abstraction records are handed to. +//! +//! The abstraction is very thin and narrow, but it allows us to avoid needing to add a +//! "build `librdkafka` with `cmake`" step to tests or local development builds. + +use std::sync::Arc; +use std::time::Duration; + +/// Sends serialized inventory records somewhere durable. +pub trait Producer { + /// What can go wrong when sending. + type Error; + + /// Enqueues one record. + /// + /// `key` controls which partition receives the message. + fn send(&self, key: &[u8], payload: Vec) -> Result<(), Self::Error>; + + /// Blocks until enqueued records have been delivered, or `timeout` elapses. + /// + /// Sending is asynchronous, so records handed over just before a process exits are + /// still sitting in a local queue. Call this during shutdown to deliver them. + fn flush(&self, timeout: Duration) -> Result<(), Self::Error>; +} + +impl Producer for Box

{ + type Error = P::Error; + + fn send(&self, key: &[u8], payload: Vec) -> Result<(), Self::Error> { + (**self).send(key, payload) + } + + fn flush(&self, timeout: Duration) -> Result<(), Self::Error> { + (**self).flush(timeout) + } +} + +impl Producer for Arc

{ + type Error = P::Error; + + fn send(&self, key: &[u8], payload: Vec) -> Result<(), Self::Error> { + (**self).send(key, payload) + } + + fn flush(&self, timeout: Duration) -> Result<(), Self::Error> { + (**self).flush(timeout) + } +} + +/// A [`Producer`] that discards everything. +#[derive(Clone, Copy, Debug, Default)] +pub struct NoopProducer; + +impl Producer for NoopProducer { + type Error = std::convert::Infallible; + + fn send(&self, _key: &[u8], _payload: Vec) -> Result<(), Self::Error> { + Ok(()) + } + + fn flush(&self, _timeout: Duration) -> Result<(), Self::Error> { + Ok(()) + } +} + +#[cfg(any(test, feature = "test-utils"))] +mod dummy { + use super::Producer; + use crate::InventoryRecord; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + /// A [`Producer`] that records everything it is given, for assertions in tests. + /// + /// Cloning shares the same buffer, so a clone can be handed to an + /// [`InventoryTracker`](crate::InventoryTracker) while the original is used to read + /// back what was emitted. + #[derive(Clone, Debug, Default)] + pub struct DummyProducer { + sent: Arc>>, + } + + /// A message as handed to [`Producer::send`], as `(key, payload)`. + pub type SentMessage = (Vec, Vec); + + impl DummyProducer { + /// Returns the raw messages sent so far, in order. + pub fn raw(&self) -> Vec { + self.sent.lock().unwrap().clone() + } + + /// Returns the records sent so far, deserialized, in order. + pub fn records(&self) -> Vec { + self.raw() + .iter() + .map(|(_, payload)| serde_json::from_slice(payload).unwrap()) + .collect() + } + + /// Discards everything sent so far. + pub fn clear(&self) { + self.sent.lock().unwrap().clear(); + } + } + + impl Producer for DummyProducer { + type Error = std::convert::Infallible; + + fn send(&self, key: &[u8], payload: Vec) -> Result<(), Self::Error> { + self.sent.lock().unwrap().push((key.to_vec(), payload)); + Ok(()) + } + + // Nothing is ever queued, so there is nothing to wait for. + fn flush(&self, _timeout: Duration) -> Result<(), Self::Error> { + Ok(()) + } + } +} + +#[cfg(any(test, feature = "test-utils"))] +pub use dummy::{DummyProducer, SentMessage}; diff --git a/objectstore-inventory-tracker/src/record.rs b/objectstore-inventory-tracker/src/record.rs new file mode 100644 index 00000000..bb3378c7 --- /dev/null +++ b/objectstore-inventory-tracker/src/record.rs @@ -0,0 +1,188 @@ +//! The wire format emitted onto the inventory topic. +//! +//! These types mirror the `shared-resources-inventory` schema registered in +//! [sentry-kafka-schemas]. The schema sets `additionalProperties: false`, so adding a +//! field here without a corresponding schema version bump produces messages that +//! consumers reject. +//! +//! [sentry-kafka-schemas]: https://github.com/getsentry/sentry-kafka-schemas + +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; + +/// The kind of change a record describes. `WRITE`, `UPDATE`, or `DELETE`. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "UPPERCASE")] +pub enum OpType { + /// The record was created, or replaced with new contents. + Write, + /// An existing record's metadata changed without its stored size changing. + /// + /// Extending a record's expiration deadline is an example of an update operation. + Update, + /// The record was removed. + Delete, +} + +/// A single inventory change event. +/// +/// Construct these through [`InventoryTracker`](crate::InventoryTracker) rather than +/// directly. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct InventoryRecord { + /// Identifies the shared resource this record belongs to. + /// + /// This is meant to match a label on a provisioned storage backend so that downstream + /// consumers can join the change stream dataset with, for instance, billing info. + pub shared_resource_id: String, + + /// The product feature this record is attributed to. + pub app_feature: String, + + /// The type of operation that occurred. + pub op_type: OpType, + + /// Opaque stable identifier, unique within `shared_resource_id`. + /// + /// A hash of the caller's storage key, derived by + /// [`InventoryTracker`](crate::InventoryTracker). The raw key is never emitted. + pub record_id: String, + + /// When the operation occurred, as Unix epoch microseconds. + pub timestamp: i64, + + /// Fraction of records the producer is emitting for this resource, in `[0, 1]`. + /// + /// `InventoryTracker` will not emit messages for records that are sampled out so the + /// consumer doesn't need to do any filtering. The reason the sample rate is included + /// on messages is so that the consumer or downstream pipelines can apply a + /// `1 / sample_rate` weight when calculating aggregates to account for changes to the + /// sample rate. + pub sample_rate: f64, + + /// Stored size in bytes. Always set for [`OpType::Write`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub size: Option, + + /// When this record is set to expire, in Unix epoch microseconds if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub expiration_time: Option, + + /// ID of the organization that owns the record, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub organization_id: Option, + + /// ID of the project that owns the record, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub project_id: Option, +} + +/// Converts a [`SystemTime`] to Unix epoch microseconds. +pub fn epoch_micros(time: SystemTime) -> i64 { + match time.duration_since(UNIX_EPOCH) { + Ok(duration) => i64::try_from(duration.as_micros()).unwrap_or(i64::MAX), + Err(err) => i64::try_from(err.duration().as_micros()).map_or(i64::MIN, |micros| -micros), + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + + use super::*; + + fn record(op_type: OpType) -> InventoryRecord { + InventoryRecord { + shared_resource_id: "example_resource".into(), + app_feature: "example_feature".into(), + op_type, + record_id: "3f7a1c2e9b4d5a6f8c0e1d2b3a4f5e6c".into(), + timestamp: 1_785_283_200_000_000, + sample_rate: 1.0, + size: None, + expiration_time: None, + organization_id: None, + project_id: None, + } + } + + #[test] + fn write_serializes_to_expected_wire_format() { + let mut rec = record(OpType::Write); + rec.size = Some(524_288); + rec.expiration_time = Some(1_785_888_000_000_000); + rec.organization_id = Some(1); + rec.project_id = Some(1); + + assert_eq!( + serde_json::to_value(&rec).unwrap(), + json!({ + "shared_resource_id": "example_resource", + "app_feature": "example_feature", + "op_type": "WRITE", + "record_id": "3f7a1c2e9b4d5a6f8c0e1d2b3a4f5e6c", + "timestamp": 1_785_283_200_000_000_i64, + "sample_rate": 1.0, + "size": 524_288, + "expiration_time": 1_785_888_000_000_000_i64, + "organization_id": 1, + "project_id": 1, + }) + ); + } + + #[test] + fn absent_optional_fields_are_omitted() { + let value = serde_json::to_value(record(OpType::Delete)).unwrap(); + let object = value.as_object().unwrap(); + + assert_eq!(object["op_type"], "DELETE"); + for absent in ["size", "expiration_time", "organization_id", "project_id"] { + assert!(!object.contains_key(absent), "{absent} should be omitted"); + } + // The required fields survive that omission. + for present in [ + "shared_resource_id", + "app_feature", + "op_type", + "record_id", + "timestamp", + "sample_rate", + ] { + assert!(object.contains_key(present), "{present} is required"); + } + } + + #[test] + fn op_types_use_uppercase_spellings() { + for (op_type, expected) in [ + (OpType::Write, "WRITE"), + (OpType::Update, "UPDATE"), + (OpType::Delete, "DELETE"), + ] { + assert_eq!(serde_json::to_value(op_type).unwrap(), json!(expected)); + } + } + + #[test] + fn epoch_micros_converts_both_directions() { + assert_eq!(epoch_micros(UNIX_EPOCH), 0); + assert_eq!( + epoch_micros(UNIX_EPOCH + Duration::from_micros(1_785_283_200_000_000)), + 1_785_283_200_000_000 + ); + assert_eq!( + epoch_micros(UNIX_EPOCH - Duration::from_micros(1_500)), + -1_500 + ); + } + + #[test] + fn epoch_micros_preserves_sub_second_precision() { + let time = UNIX_EPOCH + Duration::new(1_785_283_200, 123_456_000); + assert_eq!(epoch_micros(time), 1_785_283_200_123_456); + } +} diff --git a/objectstore-inventory-tracker/src/tracker.rs b/objectstore-inventory-tracker/src/tracker.rs new file mode 100644 index 00000000..40d1dcd2 --- /dev/null +++ b/objectstore-inventory-tracker/src/tracker.rs @@ -0,0 +1,492 @@ +//! Record identity, sampling, and the emitting entry point. + +use std::time::{Duration, SystemTime}; + +use crate::producer::Producer; +use crate::record::{InventoryRecord, OpType, epoch_micros}; + +/// Hex characters, lowercase, matching the format consumers expect in `record_id`. +const HEX: &[u8; 16] = b"0123456789abcdef"; + +/// Bytes of the hash reserved for the sampling decision. +const TOKEN_BYTES: usize = 8; +/// Bytes of the hash used to build the emitted record id, giving a 128-bit identifier. +const ID_BYTES: usize = 16; + +// Don't need a dep just for this. +fn hex_encode(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len() * 2); + for &byte in bytes { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + out +} + +/// Emits inventory records for one shared resource. +/// +/// # Hashing and sampling +/// +/// `InventoryTracker` decides whether to emit a message for a given record based on the +/// configured `sample_rate` and a hash of the record ID passed in by the caller. A sample +/// rate of 1.0 means it will emit messages for 100% of records. A sample rate of 0.5 +/// means it will emit messages for 50% of records. If a record is sampled out, no change +/// to that record will ever emit a message. If a record is included in the sample, every +/// change to that record will emit a message. +/// +/// The hash is what `InventoryTracker` actually uses to populate the `record_id` message +/// field. Each message also includes the sample rate that was in effect at the time. +/// +/// **The hash is a permanent wire contract.** Changing the algorithm, the byte ranges, +/// or introducing a salt renames every record, and consumers will double-count until the +/// old identifiers age out. +/// +/// # Example +/// +/// ``` +/// use objectstore_inventory_tracker::{InventoryTracker, NoopProducer}; +/// use std::time::SystemTime; +/// +/// let tracker = InventoryTracker::new(NoopProducer, "example_resource", 1.0); +/// +/// tracker.write( +/// "example_feature/org.123/project.456/objects/abc", +/// "example_feature", +/// 4096, +/// SystemTime::now(), +/// None, +/// Some(123), +/// Some(456), +/// )?; +/// # Ok::<(), std::convert::Infallible>(()) +/// ``` +#[derive(Clone, Debug)] +pub struct InventoryTracker { + producer: P, + shared_resource_id: String, + sample_rate: f64, + /// `sample_rate` as a point in the token space, so sampling is an integer comparison. + sample_threshold: u64, +} + +impl InventoryTracker

{ + /// Creates a tracker emitting for `shared_resource_id` at `sample_rate`. + /// + /// `shared_resource_id` is meant to match a label on a provisioned storage backend so + /// that downstream consumers can join the change stream dataset with, for instance, + /// billing info. + /// + /// `sample_rate` is clamped to `[0, 1]`. A rate of `1.0` emits every record and + /// short-circuits the sampling check entirely. A rate of `0.0` emits no records. + pub fn new(producer: P, shared_resource_id: impl Into, sample_rate: f64) -> Self { + let sample_rate = if sample_rate.is_nan() { + 1.0 + } else { + sample_rate.clamp(0.0, 1.0) + }; + + Self { + producer, + shared_resource_id: shared_resource_id.into(), + sample_rate, + // If a record's hash token is greater than or equal to this threshold, it is + // skipped. A sample rate of 0 produces a threshold of 0 so all tokens are + // skipped. A sample rate of 1.0 produces a threshold of `u64::MAX` so all + // tokens will be included (except when the token _is_ `u64::MAX`, unlikely as + // that may be. That edge case is handled in `sample()`). + sample_threshold: (sample_rate * (u64::MAX as f64)) as u64, + } + } + + /// The storage resource this tracker emits for. + pub fn shared_resource_id(&self) -> &str { + &self.shared_resource_id + } + + /// The fraction of records being emitted. + pub fn sample_rate(&self) -> f64 { + self.sample_rate + } + + /// Decides whether `storage_key` is tracked, returning its record id if so. + /// + /// The decision is deterministic for a `storage_key`, so every operation on a record + /// resolves the same way. + /// + /// It is also monotone in the rate: the set sampled at a lower rate is a subset of + /// the set sampled at a higher one. Raising the rate is therefore safe for records + /// already in flight. + fn sample(&self, storage_key: &str) -> Option { + if self.sample_threshold == 0 { + return None; + } + + let hash = blake3::hash(storage_key.as_bytes()); + let bytes = hash.as_bytes(); + + // This `sample_threshold < u64::MAX` check is handling an edge case. A sample + // rate of 1.0 produces a sample threshold of `u64::MAX` which should include + // everything. However, technically it will incorrectly skip hash tokens that + // happen to equal `u64::MAX`, however unlikely that is. So, if our threshold is + // `u64::MAX`, we skip this sampling check and just return the encoded key. + if self.sample_threshold < u64::MAX { + let token = u64::from_le_bytes(bytes[..TOKEN_BYTES].try_into().ok()?); + if token >= self.sample_threshold { + return None; + } + } + + Some(hex_encode(&bytes[TOKEN_BYTES..TOKEN_BYTES + ID_BYTES])) + } + + /// Emits a `WRITE`: the record was created, or replaced with new contents. + /// + /// Does nothing and returns `Ok(())` if `storage_key` is not sampled. + #[allow(clippy::too_many_arguments)] + pub fn write( + &self, + storage_key: &str, + app_feature: &str, + size: u64, + timestamp: SystemTime, + expiration_time: Option, + organization_id: Option, + project_id: Option, + ) -> Result<(), P::Error> { + let Some(record_id) = self.sample(storage_key) else { + return Ok(()); + }; + + self.emit(InventoryRecord { + shared_resource_id: self.shared_resource_id.clone(), + app_feature: app_feature.to_owned(), + op_type: OpType::Write, + record_id, + timestamp: epoch_micros(timestamp), + sample_rate: self.sample_rate, + size: Some(size), + expiration_time: expiration_time.map(epoch_micros), + organization_id, + project_id, + }) + } + + /// Emits an `UPDATE`: metadata changed but the stored size did not. + /// + /// Does nothing and returns `Ok(())` if `storage_key` is not sampled. + pub fn update( + &self, + storage_key: &str, + app_feature: &str, + timestamp: SystemTime, + expiration_time: Option, + organization_id: Option, + project_id: Option, + ) -> Result<(), P::Error> { + let Some(record_id) = self.sample(storage_key) else { + return Ok(()); + }; + + self.emit(InventoryRecord { + shared_resource_id: self.shared_resource_id.clone(), + app_feature: app_feature.to_owned(), + op_type: OpType::Update, + record_id, + timestamp: epoch_micros(timestamp), + sample_rate: self.sample_rate, + // Omitted, which consumers read as "unchanged" and carry forward. + size: None, + expiration_time: expiration_time.map(epoch_micros), + organization_id, + project_id, + }) + } + + /// Emits a `DELETE`: the record is gone. + /// + /// Does nothing and returns `Ok(())` if `storage_key` is not sampled. + pub fn delete( + &self, + storage_key: &str, + app_feature: &str, + timestamp: SystemTime, + ) -> Result<(), P::Error> { + let Some(record_id) = self.sample(storage_key) else { + return Ok(()); + }; + + self.emit(InventoryRecord { + shared_resource_id: self.shared_resource_id.clone(), + app_feature: app_feature.to_owned(), + op_type: OpType::Delete, + record_id, + timestamp: epoch_micros(timestamp), + sample_rate: self.sample_rate, + size: None, + expiration_time: None, + organization_id: None, + project_id: None, + }) + } + + /// Blocks until emitted records have been delivered, or `timeout` elapses. + /// + /// Call this during shutdown, within whatever budget the service allows for draining. + /// Without it, records emitted moments before exit are still in a local queue and are + /// lost with the process. + pub fn flush(&self, timeout: Duration) -> Result<(), P::Error> { + self.producer.flush(timeout) + } + + fn emit(&self, record: InventoryRecord) -> Result<(), P::Error> { + let key = record.record_id.clone(); + // Serialization of this struct cannot fail: every field is a plain scalar or + // string. + let payload = serde_json::to_vec(&record).expect("inventory record is serializable"); + self.producer.send(key.as_bytes(), payload) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use crate::producer::DummyProducer; + + use super::*; + + fn tracker(rate: f64) -> (DummyProducer, InventoryTracker) { + let producer = DummyProducer::default(); + let tracker = InventoryTracker::new(producer.clone(), "example_resource", rate); + (producer, tracker) + } + + #[test] + fn record_id_is_stable_for_a_given_key() { + let (_, tracker) = tracker(1.0); + let key = "example_feature/org.1/project.1/objects/abc"; + assert_eq!(tracker.sample(key), tracker.sample(key)); + + let record_id = tracker.sample(key).unwrap(); + assert_eq!(record_id.len(), ID_BYTES * 2); + assert!( + record_id + .chars() + .all(|c| c.is_ascii_hexdigit() && !c.is_uppercase()) + ); + } + + #[test] + fn distinct_keys_get_distinct_record_ids() { + let (_, tracker) = tracker(1.0); + let ids: HashSet<_> = (0..1000) + .map(|i| tracker.sample(&format!("key/{i}")).unwrap()) + .collect(); + assert_eq!(ids.len(), 1000); + } + + #[test] + fn rate_of_one_samples_everything() { + let (_, tracker) = tracker(1.0); + for i in 0..1000 { + assert!(tracker.sample(&format!("key/{i}")).is_some()); + } + } + + #[test] + fn sampling_is_uniform_across_key_prefixes() { + let rate = 0.25; + let (_, tracker) = tracker(rate); + + for prefix in ["attachments", "profiles", "preprod", "trace_attachments"] { + let total = 20_000; + let sampled = (0..total) + .filter(|i| { + tracker + .sample(&format!("{prefix}/org.1/project.1/objects/{i}")) + .is_some() + }) + .count(); + let observed = sampled as f64 / total as f64; + assert!( + (observed - rate).abs() < 0.02, + "prefix {prefix} sampled at {observed}, expected ~{rate}" + ); + } + } + + #[test] + fn sampling_is_monotone_in_the_rate() { + let (_, low) = tracker(0.1); + let (_, high) = tracker(0.5); + let (_, full) = tracker(1.0); + + for i in 0..5000 { + let key = format!("key/{i}"); + if low.sample(&key).is_some() { + assert!( + high.sample(&key).is_some(), + "{key} dropped when raising to 0.5" + ); + assert!( + full.sample(&key).is_some(), + "{key} dropped when raising to 1.0" + ); + } + } + } + + #[test] + fn record_id_is_independent_of_sample_rate() { + let (_, low) = tracker(0.1); + let (_, full) = tracker(1.0); + + for i in 0..2000 { + let key = format!("key/{i}"); + if let Some(sampled) = low.sample(&key) { + assert_eq!(sampled, full.sample(&key).unwrap()); + } + } + } + + #[test] + fn emitted_message_is_keyed_on_the_record_id() { + let (producer, tracker) = tracker(1.0); + tracker + .write( + "some/key", + "example_feature", + 10, + SystemTime::now(), + None, + None, + None, + ) + .unwrap(); + + let (message_key, _) = producer.raw().into_iter().next().unwrap(); + let record_id = &producer.records()[0].record_id; + assert_eq!(message_key, record_id.as_bytes()); + assert_ne!( + message_key, b"some/key", + "the raw storage key must not be used as the message key" + ); + } + + #[test] + fn write_always_carries_a_size_and_delete_never_does() { + let (producer, tracker) = tracker(1.0); + let now = SystemTime::now(); + + tracker + .write("some/key", "f", 4096, now, None, None, None) + .unwrap(); + tracker + .update("some/key", "f", now, Some(now), None, None) + .unwrap(); + tracker.delete("some/key", "f", now).unwrap(); + + let records = producer.records(); + assert_eq!(records[0].op_type, OpType::Write); + assert_eq!(records[0].size, Some(4096)); + assert_eq!(records[1].op_type, OpType::Update); + assert_eq!(records[1].size, None, "update means size unchanged"); + assert_eq!(records[2].op_type, OpType::Delete); + assert_eq!(records[2].size, None); + } + + #[test] + fn every_operation_on_a_key_reports_the_same_record_id() { + let (producer, tracker) = tracker(1.0); + let now = SystemTime::now(); + + tracker + .write("some/key", "f", 4096, now, None, None, None) + .unwrap(); + tracker.delete("some/key", "f", now).unwrap(); + + let records = producer.records(); + assert_eq!(records[0].record_id, records[1].record_id); + } + + #[test] + fn unsampled_keys_emit_nothing() { + let (producer, tracker) = tracker(0.25); + let now = SystemTime::now(); + + let unsampled = (0..) + .map(|i| format!("key/{i}")) + .find(|key| tracker.sample(key).is_none()) + .expect("some key is not sampled"); + + tracker + .write(&unsampled, "f", 1, now, None, None, None) + .unwrap(); + tracker + .update(&unsampled, "f", now, None, None, None) + .unwrap(); + tracker.delete(&unsampled, "f", now).unwrap(); + + assert!(producer.records().is_empty()); + } + + #[test] + fn sample_rate_is_stamped_on_every_record() { + let (producer, tracker) = tracker(0.25); + let sampled = (0..) + .map(|i| format!("key/{i}")) + .find(|key| tracker.sample(key).is_some()) + .expect("some key is sampled"); + + tracker + .write(&sampled, "f", 1, SystemTime::now(), None, None, None) + .unwrap(); + + assert_eq!(producer.records()[0].sample_rate, 0.25); + } + + /// The threshold is the one place float arithmetic survives, so pin what it derives + /// to — particularly at the endpoints, where the cast saturates. + #[test] + fn threshold_is_derived_from_the_rate() { + for (rate, expected) in [ + (0.0, 0), + (0.25, 1u64 << 62), + (0.5, 1u64 << 63), + (1.0, u64::MAX), + ] { + assert_eq!(tracker(rate).1.sample_threshold, expected, "rate {rate}"); + } + } + + #[test] + fn out_of_range_rates_are_clamped() { + for high in [f64::NAN, f64::INFINITY, 2.0] { + let (_, tracker) = tracker(high); + assert_eq!( + tracker.sample_rate(), + 1.0, + "rate {high} should clamp to 1.0" + ); + } + for low in [f64::NEG_INFINITY, -1.0] { + let (_, tracker) = tracker(low); + assert_eq!(tracker.sample_rate(), 0.0, "rate {low} should clamp to 0.0"); + } + } + + #[test] + fn a_rate_of_zero_emits_nothing() { + let (producer, tracker) = tracker(0.0); + let now = SystemTime::now(); + + for i in 0..1000 { + let key = format!("key/{i}"); + tracker.write(&key, "f", 1, now, None, None, None).unwrap(); + tracker.update(&key, "f", now, None, None, None).unwrap(); + tracker.delete(&key, "f", now).unwrap(); + } + + assert!(producer.records().is_empty()); + } +}