objectstore_inventory_tracker/
tracker.rs

1//! Record identity, sampling, and the emitting entry point.
2
3use std::future::Future;
4use std::time::{Duration, SystemTime};
5
6use crate::producer::Producer;
7use crate::record::{InventoryRecord, OpType, epoch_micros};
8
9/// Hex characters, lowercase, matching the format consumers expect in `record_id`.
10const HEX: &[u8; 16] = b"0123456789abcdef";
11
12/// Bytes of the hash reserved for the sampling decision.
13const TOKEN_BYTES: usize = 8;
14/// Bytes of the hash used to build the emitted record id, giving a 128-bit identifier.
15const ID_BYTES: usize = 16;
16
17// Don't need a dep just for this.
18fn hex_encode(bytes: &[u8]) -> String {
19    let mut out = String::with_capacity(bytes.len() * 2);
20    for &byte in bytes {
21        out.push(HEX[(byte >> 4) as usize] as char);
22        out.push(HEX[(byte & 0x0f) as usize] as char);
23    }
24    out
25}
26
27/// Emits inventory records for one shared resource.
28///
29/// # Hashing and sampling
30///
31/// `InventoryTracker` decides whether to emit a message for a given record based on the
32/// configured `sample_rate` and a hash of the record ID passed in by the caller. A sample
33/// rate of 1.0 means it will emit messages for 100% of records. A sample rate of 0.5
34/// means it will emit messages for 50% of records. If a record is sampled out, no change
35/// to that record will ever emit a message. If a record is included in the sample, every
36/// change to that record will emit a message.
37///
38/// The hash is what `InventoryTracker` actually uses to populate the `record_id` message
39/// field. Each message also includes the sample rate that was in effect at the time.
40///
41/// **The hash is a permanent wire contract.** Changing the algorithm, the byte ranges,
42/// or introducing a salt renames every record, and consumers will double-count until the
43/// old identifiers age out.
44///
45/// # Example
46///
47/// ```
48/// use objectstore_inventory_tracker::{InventoryTracker, NoopProducer};
49/// use std::time::SystemTime;
50///
51/// let tracker = InventoryTracker::new(NoopProducer, "example_resource", 1.0);
52///
53/// tracker.write(
54///     "example_feature/org.123/project.456/objects/abc",
55///     "example_feature",
56///     4096,
57///     SystemTime::now(),
58///     None,
59///     Some(123),
60///     Some(456),
61/// )?;
62/// # Ok::<(), std::convert::Infallible>(())
63/// ```
64#[derive(Clone, Debug)]
65pub struct InventoryTracker<P: Producer> {
66    producer: P,
67    shared_resource_id: String,
68    sample_rate: f64,
69    /// `sample_rate` as a point in the token space, so sampling is an integer comparison.
70    sample_threshold: u64,
71}
72
73impl<P: Producer> InventoryTracker<P> {
74    /// Creates a tracker emitting for `shared_resource_id` at `sample_rate`.
75    ///
76    /// `shared_resource_id` is meant to match a label on a provisioned storage backend so
77    /// that downstream consumers can join the change stream dataset with, for instance,
78    /// billing info.
79    ///
80    /// `sample_rate` is clamped to `[0, 1]`. A rate of `1.0` emits every record and
81    /// short-circuits the sampling check entirely. A rate of `0.0` emits no records.
82    pub fn new(producer: P, shared_resource_id: impl Into<String>, sample_rate: f64) -> Self {
83        let sample_rate = if sample_rate.is_nan() {
84            1.0
85        } else {
86            sample_rate.clamp(0.0, 1.0)
87        };
88
89        Self {
90            producer,
91            shared_resource_id: shared_resource_id.into(),
92            sample_rate,
93            // If a record's hash token is greater than or equal to this threshold, it is
94            // skipped. A sample rate of 0 produces a threshold of 0 so all tokens are
95            // skipped. A sample rate of 1.0 produces a threshold of `u64::MAX` so all
96            // tokens will be included (except when the token _is_ `u64::MAX`, unlikely as
97            // that may be. That edge case is handled in `sample()`).
98            sample_threshold: (sample_rate * (u64::MAX as f64)) as u64,
99        }
100    }
101
102    /// The storage resource this tracker emits for.
103    pub fn shared_resource_id(&self) -> &str {
104        &self.shared_resource_id
105    }
106
107    /// The fraction of records being emitted.
108    pub fn sample_rate(&self) -> f64 {
109        self.sample_rate
110    }
111
112    /// Decides whether `storage_key` is tracked, returning its record id if so.
113    ///
114    /// The decision is deterministic for a `storage_key`, so every operation on a record
115    /// resolves the same way.
116    ///
117    /// It is also monotone in the rate: the set sampled at a lower rate is a subset of
118    /// the set sampled at a higher one. Raising the rate is therefore safe for records
119    /// already in flight.
120    fn sample(&self, storage_key: &str) -> Option<String> {
121        if self.sample_threshold == 0 {
122            return None;
123        }
124
125        let hash = blake3::hash(storage_key.as_bytes());
126        let bytes = hash.as_bytes();
127
128        // This `sample_threshold < u64::MAX` check is handling an edge case. A sample
129        // rate of 1.0 produces a sample threshold of `u64::MAX` which should include
130        // everything. However, technically it will incorrectly skip hash tokens that
131        // happen to equal `u64::MAX`, however unlikely that is. So, if our threshold is
132        // `u64::MAX`, we skip this sampling check and just return the encoded key.
133        if self.sample_threshold < u64::MAX {
134            let token = u64::from_le_bytes(bytes[..TOKEN_BYTES].try_into().ok()?);
135            if token >= self.sample_threshold {
136                return None;
137            }
138        }
139
140        Some(hex_encode(&bytes[TOKEN_BYTES..TOKEN_BYTES + ID_BYTES]))
141    }
142
143    /// Emits a `WRITE`: the record was created, or replaced with new contents.
144    ///
145    /// Does nothing and returns `Ok(())` if `storage_key` is not sampled.
146    #[allow(clippy::too_many_arguments)]
147    pub fn write(
148        &self,
149        storage_key: &str,
150        app_feature: &str,
151        size: u64,
152        timestamp: SystemTime,
153        expiration_time: Option<SystemTime>,
154        organization_id: Option<u64>,
155        project_id: Option<u64>,
156    ) -> Result<(), P::Error> {
157        let Some(record_id) = self.sample(storage_key) else {
158            return Ok(());
159        };
160
161        self.emit(InventoryRecord {
162            shared_resource_id: self.shared_resource_id.clone(),
163            app_feature: app_feature.to_owned(),
164            op_type: OpType::Write,
165            record_id,
166            timestamp: epoch_micros(timestamp),
167            sample_rate: self.sample_rate,
168            size: Some(size),
169            expiration_time: expiration_time.map(epoch_micros),
170            organization_id,
171            project_id,
172        })
173    }
174
175    /// Emits an `UPDATE`: metadata changed but the stored size did not.
176    ///
177    /// Does nothing and returns `Ok(())` if `storage_key` is not sampled.
178    pub fn update(
179        &self,
180        storage_key: &str,
181        app_feature: &str,
182        timestamp: SystemTime,
183        expiration_time: Option<SystemTime>,
184        organization_id: Option<u64>,
185        project_id: Option<u64>,
186    ) -> Result<(), P::Error> {
187        let Some(record_id) = self.sample(storage_key) else {
188            return Ok(());
189        };
190
191        self.emit(InventoryRecord {
192            shared_resource_id: self.shared_resource_id.clone(),
193            app_feature: app_feature.to_owned(),
194            op_type: OpType::Update,
195            record_id,
196            timestamp: epoch_micros(timestamp),
197            sample_rate: self.sample_rate,
198            // Omitted, which consumers read as "unchanged" and carry forward.
199            size: None,
200            expiration_time: expiration_time.map(epoch_micros),
201            organization_id,
202            project_id,
203        })
204    }
205
206    /// Emits a `DELETE`: the record is gone.
207    ///
208    /// Does nothing and returns `Ok(())` if `storage_key` is not sampled.
209    pub fn delete(
210        &self,
211        storage_key: &str,
212        app_feature: &str,
213        timestamp: SystemTime,
214    ) -> Result<(), P::Error> {
215        let Some(record_id) = self.sample(storage_key) else {
216            return Ok(());
217        };
218
219        self.emit(InventoryRecord {
220            shared_resource_id: self.shared_resource_id.clone(),
221            app_feature: app_feature.to_owned(),
222            op_type: OpType::Delete,
223            record_id,
224            timestamp: epoch_micros(timestamp),
225            sample_rate: self.sample_rate,
226            size: None,
227            expiration_time: None,
228            organization_id: None,
229            project_id: None,
230        })
231    }
232
233    /// Waits for emitted records to be delivered, or until `timeout` elapses.
234    ///
235    /// Call this during shutdown, within whatever budget the service allows for draining.
236    /// Without it, records emitted moments before exit are still in a local queue and are
237    /// lost with the process.
238    pub fn join(
239        &self,
240        timeout: Duration,
241    ) -> impl Future<Output = Result<(), P::Error>> + Send + use<P>
242    where
243        P: Clone + Send + Sync + 'static,
244        P::Error: Send + 'static,
245    {
246        self.producer.join(timeout)
247    }
248
249    fn emit(&self, record: InventoryRecord) -> Result<(), P::Error> {
250        let key = record.record_id.clone();
251        // Serialization of this struct cannot fail: every field is a plain scalar or
252        // string.
253        let payload = serde_json::to_vec(&record).expect("inventory record is serializable");
254        self.producer.send(key.as_bytes(), payload)
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use std::collections::HashSet;
261
262    use crate::test_utils::DummyProducer;
263
264    use super::*;
265
266    fn tracker(rate: f64) -> (DummyProducer, InventoryTracker<DummyProducer>) {
267        let producer = DummyProducer::default();
268        let tracker = InventoryTracker::new(producer.clone(), "example_resource", rate);
269        (producer, tracker)
270    }
271
272    #[test]
273    fn record_id_is_stable_for_a_given_key() {
274        let (_, tracker) = tracker(1.0);
275        let key = "example_feature/org.1/project.1/objects/abc";
276        assert_eq!(tracker.sample(key), tracker.sample(key));
277
278        let record_id = tracker.sample(key).unwrap();
279        assert_eq!(record_id.len(), ID_BYTES * 2);
280        assert!(
281            record_id
282                .chars()
283                .all(|c| c.is_ascii_hexdigit() && !c.is_uppercase())
284        );
285    }
286
287    #[test]
288    fn distinct_keys_get_distinct_record_ids() {
289        let (_, tracker) = tracker(1.0);
290        let ids: HashSet<_> = (0..1000)
291            .map(|i| tracker.sample(&format!("key/{i}")).unwrap())
292            .collect();
293        assert_eq!(ids.len(), 1000);
294    }
295
296    #[test]
297    fn rate_of_one_samples_everything() {
298        let (_, tracker) = tracker(1.0);
299        for i in 0..1000 {
300            assert!(tracker.sample(&format!("key/{i}")).is_some());
301        }
302    }
303
304    #[test]
305    fn sampling_is_uniform_across_key_prefixes() {
306        let rate = 0.25;
307        let (_, tracker) = tracker(rate);
308
309        for prefix in ["attachments", "profiles", "preprod", "trace_attachments"] {
310            let total = 20_000;
311            let sampled = (0..total)
312                .filter(|i| {
313                    tracker
314                        .sample(&format!("{prefix}/org.1/project.1/objects/{i}"))
315                        .is_some()
316                })
317                .count();
318            let observed = sampled as f64 / total as f64;
319            assert!(
320                (observed - rate).abs() < 0.02,
321                "prefix {prefix} sampled at {observed}, expected ~{rate}"
322            );
323        }
324    }
325
326    #[test]
327    fn sampling_is_monotone_in_the_rate() {
328        let (_, low) = tracker(0.1);
329        let (_, high) = tracker(0.5);
330        let (_, full) = tracker(1.0);
331
332        for i in 0..5000 {
333            let key = format!("key/{i}");
334            if low.sample(&key).is_some() {
335                assert!(
336                    high.sample(&key).is_some(),
337                    "{key} dropped when raising to 0.5"
338                );
339                assert!(
340                    full.sample(&key).is_some(),
341                    "{key} dropped when raising to 1.0"
342                );
343            }
344        }
345    }
346
347    #[test]
348    fn record_id_is_independent_of_sample_rate() {
349        let (_, low) = tracker(0.1);
350        let (_, full) = tracker(1.0);
351
352        for i in 0..2000 {
353            let key = format!("key/{i}");
354            if let Some(sampled) = low.sample(&key) {
355                assert_eq!(sampled, full.sample(&key).unwrap());
356            }
357        }
358    }
359
360    #[test]
361    fn emitted_message_is_keyed_on_the_record_id() {
362        let (producer, tracker) = tracker(1.0);
363        tracker
364            .write(
365                "some/key",
366                "example_feature",
367                10,
368                SystemTime::now(),
369                None,
370                None,
371                None,
372            )
373            .unwrap();
374
375        let (message_key, _) = producer.raw().into_iter().next().unwrap();
376        let record_id = &producer.records()[0].record_id;
377        assert_eq!(message_key, record_id.as_bytes());
378        assert_ne!(
379            message_key, b"some/key",
380            "the raw storage key must not be used as the message key"
381        );
382    }
383
384    #[test]
385    fn write_always_carries_a_size_and_delete_never_does() {
386        let (producer, tracker) = tracker(1.0);
387        let now = SystemTime::now();
388
389        tracker
390            .write("some/key", "f", 4096, now, None, None, None)
391            .unwrap();
392        tracker
393            .update("some/key", "f", now, Some(now), None, None)
394            .unwrap();
395        tracker.delete("some/key", "f", now).unwrap();
396
397        let records = producer.records();
398        assert_eq!(records[0].op_type, OpType::Write);
399        assert_eq!(records[0].size, Some(4096));
400        assert_eq!(records[1].op_type, OpType::Update);
401        assert_eq!(records[1].size, None, "update means size unchanged");
402        assert_eq!(records[2].op_type, OpType::Delete);
403        assert_eq!(records[2].size, None);
404    }
405
406    #[test]
407    fn every_operation_on_a_key_reports_the_same_record_id() {
408        let (producer, tracker) = tracker(1.0);
409        let now = SystemTime::now();
410
411        tracker
412            .write("some/key", "f", 4096, now, None, None, None)
413            .unwrap();
414        tracker.delete("some/key", "f", now).unwrap();
415
416        let records = producer.records();
417        assert_eq!(records[0].record_id, records[1].record_id);
418    }
419
420    #[test]
421    fn unsampled_keys_emit_nothing() {
422        let (producer, tracker) = tracker(0.25);
423        let now = SystemTime::now();
424
425        let unsampled = (0..)
426            .map(|i| format!("key/{i}"))
427            .find(|key| tracker.sample(key).is_none())
428            .expect("some key is not sampled");
429
430        tracker
431            .write(&unsampled, "f", 1, now, None, None, None)
432            .unwrap();
433        tracker
434            .update(&unsampled, "f", now, None, None, None)
435            .unwrap();
436        tracker.delete(&unsampled, "f", now).unwrap();
437
438        assert!(producer.records().is_empty());
439    }
440
441    #[test]
442    fn sample_rate_is_stamped_on_every_record() {
443        let (producer, tracker) = tracker(0.25);
444        let sampled = (0..)
445            .map(|i| format!("key/{i}"))
446            .find(|key| tracker.sample(key).is_some())
447            .expect("some key is sampled");
448
449        tracker
450            .write(&sampled, "f", 1, SystemTime::now(), None, None, None)
451            .unwrap();
452
453        assert_eq!(producer.records()[0].sample_rate, 0.25);
454    }
455
456    #[test]
457    fn threshold_is_derived_from_the_rate() {
458        for (rate, expected) in [
459            (0.0, 0),
460            (0.25, 1u64 << 62),
461            (0.5, 1u64 << 63),
462            (1.0, u64::MAX),
463        ] {
464            assert_eq!(tracker(rate).1.sample_threshold, expected, "rate {rate}");
465        }
466    }
467
468    #[test]
469    fn out_of_range_rates_are_clamped() {
470        for high in [f64::NAN, f64::INFINITY, 2.0] {
471            let (_, tracker) = tracker(high);
472            assert_eq!(
473                tracker.sample_rate(),
474                1.0,
475                "rate {high} should clamp to 1.0"
476            );
477        }
478        for low in [f64::NEG_INFINITY, -1.0] {
479            let (_, tracker) = tracker(low);
480            assert_eq!(tracker.sample_rate(), 0.0, "rate {low} should clamp to 0.0");
481        }
482    }
483
484    #[test]
485    fn a_rate_of_zero_emits_nothing() {
486        let (producer, tracker) = tracker(0.0);
487        let now = SystemTime::now();
488
489        for i in 0..1000 {
490            let key = format!("key/{i}");
491            tracker.write(&key, "f", 1, now, None, None, None).unwrap();
492            tracker.update(&key, "f", now, None, None, None).unwrap();
493            tracker.delete(&key, "f", now).unwrap();
494        }
495
496        assert!(producer.records().is_empty());
497    }
498}
objectstore_inventory_tracker/tracker.rs

objectstore_inventory_tracker/
tracker.rs
