add atomic flush protocol and tsdb

Author: agavra

open-tsdb/src/db.rs

@@ -0,0 +1,121 @@
+// Tsdb is the coordination layer for OpenTSDB, ensuring the ingestion and query
+// layers are synchronized. It maintains a three-level hierarchy:
+//
+// 1. Mutable Head: this stores incoming data that has not yet been
+//    flushed to storage. They are updated via deltas, which are applied to
+//    the ehad chunk for the corresponding time bucket.
+//
+// 2. Frozen Head: on a trigger (currently time-based), the head
+//    is frozen and its data is ready to be flushed to storage. Since the time
+//    it takes to flush the data is variable, the frozen head chunks are maintained
+//    in memory until the flush is complete. At this point the frozen chunk is
+//    atomically discarded with the update of the storage snapshot.
+//
+// 3. Storage Snapshot: a snapshot of the storage layer that does not yet have the
+//    any of the head chunks' data. This is used for consistency as it is possible
+//    that data has already been flushed to storage before the frozen chunks are
+//    discarded (which would result in duplicate reads).
+//
+// In addition to the three levels, the Tsdb maintains a cache of the entire series
+// dictionary for each active time bucket since this is required for ingestion.
+
+use std::sync::{Arc, atomic::AtomicU32};
+
+use dashmap::DashMap;
+use opendata_common::{Storage, storage::StorageSnapshot};
+use tokio::sync::{Mutex, RwLock};
+
+use crate::delta::TsdbDelta;
+use crate::storage::OpenTsdbStorageReadExt;
+use crate::util::Result;
+use crate::{
+    head::TsdbHead,
+    model::{SeriesFingerprint, SeriesId, TimeBucket},
+};
+
+pub(crate) struct TsdbState {
+    head: TsdbHead,
+    frozen_head: Option<TsdbHead>,
+    storage_snapshot: Arc<dyn StorageSnapshot>,
+}
+
+pub(crate) struct Tsdb {
+    bucket: TimeBucket,
+    series_dict: DashMap<SeriesFingerprint, SeriesId>,
+    next_series_id: AtomicU32,
+    /// The state of the Tsdb is protected by a read-write
+    /// lock to allow us to atomically flush and update the
+    /// storage snapshot. Note that only the read lock is
+    /// required for queries and ingestion (which modifies the
+    /// underlying data structures in a thread-safe way)
+    state: Arc<RwLock<TsdbState>>,
+    /// Mutex to ensure only one flush operation can run at a time.
+    /// This prevents concurrent flushes from interfering with each other.
+    flush_mutex: Arc<Mutex<()>>,
+}
+
+impl Tsdb {
+    pub(crate) async fn load(bucket: TimeBucket, storage: Arc<dyn Storage>) -> Result<Self> {
+        let series_dict = DashMap::new();
+        let next_series_id = storage
+            .load_series_dictionary(&bucket, |fingerprint, series_id| {
+                series_dict.insert(fingerprint, series_id);
+            })
+            .await?;
+
+        Ok(Self {
+            bucket: bucket.clone(),
+            series_dict,
+            next_series_id: AtomicU32::new(next_series_id),
+            state: Arc::new(RwLock::new(TsdbState {
+                head: TsdbHead::new(bucket.clone()),
+                frozen_head: None,
+                storage_snapshot: storage.snapshot().await?,
+            })),
+            flush_mutex: Arc::new(Mutex::new(())),
+        })
+    }
+
+    pub(crate) async fn ingest(&self, delta: TsdbDelta) -> Result<()> {
+        let state = self.state.read().await;
+        state.head.merge(&delta)?;
+        Ok(())
+    }
+
+    pub(crate) async fn flush(&self, storage: Arc<dyn Storage>) -> Result<()> {
+        let _flush_guard = self.flush_mutex.lock().await;
+
+        // blocking section: freeze the head and replaces it
+        // with a new head block, keeping a reference to the old
+        // frozen head
+        {
+            let mut state = self.state.write().await;
+            state.head.freeze();
+            let frozen_head =
+                std::mem::replace(&mut state.head, TsdbHead::new(self.bucket.clone()));
+            state.frozen_head = Some(frozen_head);
+        }
+
+        // non-blocking section: flush the frozen head to storage
+        // this can take time so its important that it only holds
+        // the read lock
+        let snapshot = {
+            let state = self.state.read().await;
+            let frozen_head = state
+                .frozen_head
+                .as_ref()
+                .expect("frozen_head should be set after write lock above");
+            // Clone storage to avoid moving it and for clarity
+            frozen_head.flush(storage.clone()).await?
+        };
+
+        // blocking section: update the storage snapshot and
+        // discard the frozen head
+        {
+            let mut state = self.state.write().await;
+            state.storage_snapshot = snapshot;
+            state.frozen_head = None;
+        }
+        Ok(())
+    }
+}

open-tsdb/src/delta.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use dashmap::DashMap;
 use opentelemetry_proto::tonic::metrics::v1::Metric;
 
 use crate::{
@@ -15,7 +16,7 @@ pub(crate) struct TsdbDeltaBuilder<'a> {
     pub(crate) bucket: TimeBucket,
     pub(crate) forward_index: ForwardIndex,
     pub(crate) inverted_index: InvertedIndex,
-    pub(crate) series_dict: &'a HashMap<SeriesFingerprint, SeriesId>,
+    pub(crate) series_dict: &'a DashMap<SeriesFingerprint, SeriesId>,
     pub(crate) series_dict_delta: HashMap<SeriesFingerprint, SeriesId>,
     pub(crate) samples: HashMap<SeriesId, Vec<Sample>>,
     pub(crate) next_series_id: u32,
@@ -24,7 +25,8 @@ pub(crate) struct TsdbDeltaBuilder<'a> {
 impl<'a> TsdbDeltaBuilder<'a> {
     pub(crate) fn new(
         bucket: TimeBucket,
-        series_dict: &'a HashMap<SeriesFingerprint, SeriesId>,
+        series_dict: &'a DashMap<SeriesFingerprint, SeriesId>,
+        next_series_id: u32,
     ) -> Self {
         Self {
             bucket,
@@ -33,7 +35,7 @@ impl<'a> TsdbDeltaBuilder<'a> {
             series_dict,
             series_dict_delta: HashMap::new(),
             samples: HashMap::new(),
-            next_series_id: 0,
+            next_series_id,
         }
     }
 
@@ -73,8 +75,8 @@ impl<'a> TsdbDeltaBuilder<'a> {
         let series_id = self
             .series_dict
             .get(&fingerprint)
-            .or_else(|| self.series_dict_delta.get(&fingerprint))
-            .copied()
+            .map(|r| *r.value())
+            .or_else(|| self.series_dict_delta.get(&fingerprint).copied())
             .unwrap_or_else(|| {
                 let series_id = self.next_series_id;
                 self.next_series_id = series_id + 1;
@@ -128,6 +130,7 @@ pub(crate) struct TsdbDelta {
 mod tests {
     use super::*;
     use crate::model::{Attribute, MetricType, Temporality};
+    use dashmap::DashMap;
 
     fn create_test_bucket() -> TimeBucket {
         TimeBucket::hour(1000)
@@ -157,8 +160,8 @@ mod tests {
     fn should_create_new_series_when_ingesting_first_sample() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = create_test_attributes();
         let sample = create_test_sample();
         let metric_unit = Some("bytes".to_string());
@@ -205,8 +208,8 @@ mod tests {
     fn should_reuse_series_id_for_samples_with_same_attributes() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = create_test_attributes();
         let sample1 = Sample {
             timestamp: 1000,
@@ -248,8 +251,8 @@ mod tests {
     fn should_create_different_series_id_for_different_attributes() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes1 = vec![Attribute {
             key: "service".to_string(),
             value: "api".to_string(),
@@ -286,13 +289,13 @@ mod tests {
     fn should_reuse_series_id_from_existing_series_dict() {
         // given
         let bucket = create_test_bucket();
-        let mut series_dict = HashMap::new();
+        let series_dict = DashMap::new();
         let mut attributes = create_test_attributes();
         // Sort attributes to match what ingest_sample does
         attributes.sort_by(|a, b| a.key.cmp(&b.key));
         let fingerprint = attributes.fingerprint();
         series_dict.insert(fingerprint, 42); // Existing series_id
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let metric_type = MetricType::Gauge;
 
         // when
@@ -314,8 +317,8 @@ mod tests {
     fn should_reuse_series_id_from_delta_dict_when_ingesting_again() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = create_test_attributes();
         let metric_type = MetricType::Gauge;
 
@@ -345,8 +348,8 @@ mod tests {
     fn should_sort_attributes_before_fingerprinting() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes1 = vec![
             Attribute {
                 key: "z_key".to_string(),
@@ -393,8 +396,8 @@ mod tests {
     fn should_store_metric_unit_and_type_in_forward_index() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = create_test_attributes();
         let metric_unit = Some("requests_per_second".to_string());
         let metric_type = MetricType::Sum {
@@ -429,8 +432,8 @@ mod tests {
     fn should_index_all_attributes_in_inverted_index() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = vec![
             Attribute {
                 key: "service".to_string(),
@@ -468,8 +471,8 @@ mod tests {
     fn should_handle_empty_attributes_list() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = Vec::<Attribute>::new();
         let metric_type = MetricType::Gauge;
 
@@ -494,8 +497,8 @@ mod tests {
     fn should_handle_none_metric_unit() {
         // given
         let bucket = create_test_bucket();
-        let series_dict = HashMap::new();
-        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict);
+        let series_dict = DashMap::new();
+        let mut builder = TsdbDeltaBuilder::new(bucket, &series_dict, 0);
         let attributes = create_test_attributes();
         let metric_type = MetricType::Gauge;
 

open-tsdb/src/tsdb.rs open-tsdb/src/head.rsopen-tsdb/src/tsdb.rs renamed to open-tsdb/src/head.rs

@@ -42,6 +42,16 @@ impl TsdbHead {
         }
     }
 
+    /// Returns a reference to the series dictionary (contains only NEW series)
+    pub fn series_dict(&self) -> &DashMap<SeriesFingerprint, SeriesId> {
+        &self.series_dict
+    }
+
+    /// Returns a reference to the bucket
+    pub fn bucket(&self) -> &TimeBucket {
+        &self.bucket
+    }
+
     pub fn merge(&self, delta: &TsdbDelta) -> Result<()> {
         if self.frozen.load(Ordering::SeqCst) {
             return Err(OpenTsdbError::Internal("TsdbHead is frozen".to_string()));
@@ -68,7 +78,7 @@ impl TsdbHead {
         self.frozen.fetch_or(true, Ordering::SeqCst)
     }
 
-    pub async fn flush(&self, storage: &dyn Storage) -> Result<Arc<dyn StorageSnapshot>> {
+    pub async fn flush(&self, storage: Arc<dyn Storage>) -> Result<Arc<dyn StorageSnapshot>> {
         if !self.frozen.load(Ordering::SeqCst) {
             return Err(OpenTsdbError::Internal(
                 "Should only flush frozen TsdbHead".to_string(),

open-tsdb/src/main.rs

@@ -1,12 +1,13 @@
 #![allow(dead_code)]
+mod db;
 mod delta;
+mod head;
 mod index;
 mod model;
 mod otel;
 mod promql;
 mod serde;
 mod storage;
-mod tsdb;
 mod util;
 
 fn main() {

open-tsdb/src/model.rs

@@ -50,7 +50,7 @@ pub(crate) enum MetricType {
     Summary,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct TimeBucket {
     pub(crate) start: BucketStart,
     pub(crate) size: BucketSize,

opendata-common/src/clock.rs

@@ -27,6 +27,12 @@ impl Clock for MockClock {
     }
 }
 
+impl Default for MockClock {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl MockClock {
     pub fn with_time(time: SystemTime) -> Self {
         Self {