TLA  Line data    Source code

       1                 : //!
       2                 : //! Shared code for consumption metics collection
       3                 : //!
       4                 : use chrono::{DateTime, Utc};
       5                 : use rand::Rng;
       6                 : use serde::{Deserialize, Serialize};
       7                 : 
       8 CBC         338 : #[derive(Serialize, serde::Deserialize, Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
       9                 : #[serde(tag = "type")]
      10                 : pub enum EventType {
      11                 :     #[serde(rename = "absolute")]
      12                 :     Absolute { time: DateTime<Utc> },
      13                 :     #[serde(rename = "incremental")]
      14                 :     Incremental {
      15                 :         start_time: DateTime<Utc>,
      16                 :         stop_time: DateTime<Utc>,
      17                 :     },
      18                 : }
      19                 : 
      20                 : impl EventType {
      21              12 :     pub fn absolute_time(&self) -> Option<&DateTime<Utc>> {
      22              12 :         use EventType::*;
      23              12 :         match self {
      24              12 :             Absolute { time } => Some(time),
      25 UBC           0 :             _ => None,
      26                 :         }
      27 CBC          12 :     }
      28                 : 
      29              11 :     pub fn incremental_timerange(&self) -> Option<std::ops::Range<&DateTime<Utc>>> {
      30              11 :         // these can most likely be thought of as Range or RangeFull, at least pageserver creates
      31              11 :         // incremental ranges where the stop and next start are equal.
      32              11 :         use EventType::*;
      33              11 :         match self {
      34                 :             Incremental {
      35              11 :                 start_time,
      36              11 :                 stop_time,
      37              11 :             } => Some(start_time..stop_time),
      38 UBC           0 :             _ => None,
      39                 :         }
      40 CBC          11 :     }
      41                 : 
      42 UBC           0 :     pub fn is_incremental(&self) -> bool {
      43               0 :         matches!(self, EventType::Incremental { .. })
      44               0 :     }
      45                 : 
      46                 :     /// Returns the absolute time, or for incremental ranges, the stop time.
      47 CBC           3 :     pub fn recorded_at(&self) -> &DateTime<Utc> {
      48               3 :         use EventType::*;
      49               3 : 
      50               3 :         match self {
      51               1 :             Absolute { time } => time,
      52               2 :             Incremental { stop_time, .. } => stop_time,
      53                 :         }
      54               3 :     }
      55                 : }
      56                 : 
      57             408 : #[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
      58                 : pub struct Event<Extra, Metric> {
      59                 :     #[serde(flatten)]
      60                 :     #[serde(rename = "type")]
      61                 :     pub kind: EventType,
      62                 : 
      63                 :     pub metric: Metric,
      64                 :     pub idempotency_key: String,
      65                 :     pub value: u64,
      66                 : 
      67                 :     #[serde(flatten)]
      68                 :     pub extra: Extra,
      69                 : }
      70                 : 
      71               4 : pub fn idempotency_key(node_id: &str) -> String {
      72               4 :     IdempotencyKey::generate(node_id).to_string()
      73               4 : }
      74                 : 
      75                 : /// Downstream users will use these to detect upload retries.
      76                 : pub struct IdempotencyKey<'a> {
      77                 :     now: chrono::DateTime<Utc>,
      78                 :     node_id: &'a str,
      79                 :     nonce: u16,
      80                 : }
      81                 : 
      82                 : impl std::fmt::Display for IdempotencyKey<'_> {
      83             119 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      84             119 :         write!(f, "{}-{}-{:04}", self.now, self.node_id, self.nonce)
      85             119 :     }
      86                 : }
      87                 : 
      88                 : impl<'a> IdempotencyKey<'a> {
      89              77 :     pub fn generate(node_id: &'a str) -> Self {
      90              77 :         IdempotencyKey {
      91              77 :             now: Utc::now(),
      92              77 :             node_id,
      93              77 :             nonce: rand::thread_rng().gen_range(0..=9999),
      94              77 :         }
      95              77 :     }
      96                 : 
      97              37 :     pub fn for_tests(now: DateTime<Utc>, node_id: &'a str, nonce: u16) -> Self {
      98              37 :         IdempotencyKey {
      99              37 :             now,
     100              37 :             node_id,
     101              37 :             nonce,
     102              37 :         }
     103              37 :     }
     104                 : }
     105                 : 
     106                 : pub const CHUNK_SIZE: usize = 1000;
     107                 : 
     108                 : // Just a wrapper around a slice of events
     109                 : // to serialize it as `{"events" : [ ] }
     110              35 : #[derive(serde::Serialize, serde::Deserialize)]
     111                 : pub struct EventChunk<'a, T: Clone> {
     112                 :     pub events: std::borrow::Cow<'a, [T]>,
     113                 : }