store segment timestamp in SegmentKey

Author: MarinPostma

libsql-wal/src/storage/backend/s3.rs

@@ -14,8 +14,8 @@ use aws_sdk_s3::operation::get_object::GetObjectOutput;
 use aws_sdk_s3::primitives::{ByteStream, SdkBody};
 use aws_sdk_s3::types::CreateBucketConfiguration;
 use aws_sdk_s3::Client;
-use aws_smithy_types_convert::date_time::DateTimeExt;
 use bytes::{Bytes, BytesMut};
+use chrono::{DateTime, Utc};
 use http_body::{Frame as HttpFrame, SizeHint};
 use libsql_sys::name::NamespaceName;
 use roaring::RoaringBitmap;
@@ -229,72 +229,9 @@ impl<IO: Io> S3Backend<IO> {
         Ok(key)
     }
 
-    /// finding a segment by timestamp should not be a common operation, and it's costly. We will
-    /// perform a binary search until we find a satisfying segment. This forces us to iteratively:
-    /// - find a segment for some frame_no
-    /// - fetch that segment index to retrieve the metata (we don't need to download the data, just
-    /// the metata is of use to us)
-    /// - check the timestamp, rince and repeat
-    ///
-    /// unfortunately, s3 doesn't have a way to list object with their metadata...
-    #[tracing::instrument(skip(self, config, folder_key))]
-    async fn find_segment_by_timestamp(&self, config: &S3Config, folder_key: &FolderKey<'_>, timestamp: DateTime<Utc>) -> Result<Option<SegmentKey>> {
-        let mut attempted_frame_no = u64::MAX;
-        let mut max_seen = 0;
-        let mut best_so_far = None;
-        loop {
-            let Some(seg_key) = self.find_segment_by_frame_no(config, folder_key, attempted_frame_no).await? else { return Ok(None) };
-            max_seen = max_seen.max(seg_key.end_frame_no);
-            let key = s3_segment_index_key(folder_key, &seg_key);
-            let object = self.s3_get(config, &key).await?;
-            let Some(segment_timestamp_rfc3339) = object.metadata().and_then(|m| m.get(SEGMENT_TIMESTAMP_META_KEY)) else { 
-                tracing::warn!("{key} is missing metadata. Cannot be used as a restore point!");
-                todo!("update next attempt frame_no");
-                // continue
-            };
-
-            let segment_timestamp = match DateTime::parse_from_rfc3339(segment_timestamp_rfc3339) {
-                Ok(s) => s.to_utc(),
-                Err(e) => {
-                    tracing::warn!("timestamp for `{key}` is invalid, skipping restore point: {e}");
-                    todo!("update next attempt frame_no");
-                    // continue
-                }
-            };
-
-            if segment_timestamp >= timestamp {
-                if segment_timestamp == timestamp {
-                    // look no further!
-                    tracing::debug!("found exact match for timestamp: {seg_key}");
-                    return Ok(Some(seg_key))
-                }
-                best_so_far = match best_so_far {
-                    Some((best, key)) => {
-                        assert!(best >= timestamp);
-                        if best - timestamp <= segment_timestamp - timestamp {
-                            Some((best, key))
-                        } else {
-                            Some((segment_timestamp, seg_key))
-                        }
-                    }
-                    None => Some((segment_timestamp, seg_key)),
-                };
-
-                let next_attempt = attempted_frame_no + ((max_seen - attempted_frame_no) / 2);
-                if next_attempt == attempted_frame_no {
-                    let (ts, key) = best_so_far.unwrap();
-                    tracing::debug!("found best match for {timestamp}: {key}@{ts}");
-                    return Ok(Some(key))
-                }
-
-                attempted_frame_no = next_attempt;
-            } else {
-                let next_attempt = attempted_frame_no - (attempted_frame_no / 2);
-                if next_attempt <= 1 {
-                    return Ok(best_so_far.map(|x| x.1))
-                }
-            }
-        }
+    // #[tracing::instrument(skip(self, config, folder_key))]
+    async fn find_segment_by_timestamp(&self, _config: &S3Config, _folder_key: &FolderKey<'_>, _timestamp: DateTime<Utc>) -> Result<Option<SegmentKey>> {
+        todo!()
     }
 
     // This method could probably be optimized a lot by using indexes and only downloading useful
@@ -407,7 +344,6 @@ impl<IO: Io> S3Backend<IO> {
                     let infos = SegmentInfo {
                         key,
                         size: entry.size().unwrap_or(0) as usize,
-                        created_at: entry.last_modified().unwrap().to_chrono_utc().unwrap(),
                     };
 
                     yield infos;

libsql-wal/src/storage/compaction/mod.rs

@@ -4,7 +4,6 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use chrono::DateTime;
-use chrono::Utc;
 use fst::map::OpBuilder;
 use fst::Streamer;
 use libsql_sys::name::NamespaceName;
@@ -58,7 +57,7 @@ impl<B> Compactor<B> {
             "CREATE TABLE IF NOT EXISTS segments (
                         start_frame_no INTEGER,
                         end_frame_no INTEGER,
-                        created_at DATE,
+                        timestamp DATE,
                         size INTEGER,
                         namespace_id INTEGER,
                         PRIMARY KEY (start_frame_no, end_frame_no),
@@ -93,7 +92,7 @@ impl<B> Compactor<B> {
     pub fn analyze(&self, namespace: &NamespaceName) -> Result<AnalyzedSegments> {
         let mut stmt = self.meta.prepare_cached(
             r#"
-        SELECT start_frame_no, end_frame_no
+        SELECT start_frame_no, end_frame_no, timestamp
         FROM segments as s
         JOIN monitored_namespaces as m
         ON m.id = s.namespace_id
@@ -105,11 +104,12 @@ impl<B> Compactor<B> {
         while let Some(row) = rows.next()? {
             let start_frame_no: u64 = row.get(0)?;
             let end_frame_no: u64 = row.get(1)?;
+            let timestamp: u64 = row.get(2)?;
             // it's free to go from one end of a segment to the next
             graph.add_edge(start_frame_no, end_frame_no, 0);
             if start_frame_no != 1 {
                 // going from a segment to the next costs us
-                graph.add_edge(start_frame_no - 1, start_frame_no, 1);
+                graph.add_edge(start_frame_no - 1, start_frame_no, timestamp);
             }
             last_frame_no = last_frame_no.max(end_frame_no);
         }
@@ -304,7 +304,7 @@ impl<B> Compactor<B> {
                     segment_id: Uuid::new_v4(),
                     start_frame_no: start,
                     end_frame_no: end,
-                    segment_timestamp: Utc::now(),
+                    segment_timestamp: DateTime::from_timestamp_millis(set.last().unwrap().timestamp as _).unwrap().to_utc(),
                 },
                 out_file,
                 out_index.into_inner().unwrap(),
@@ -375,9 +375,13 @@ impl AnalyzedSegments {
         match path {
             Some((_len, nodes)) => {
                 for chunk in nodes.chunks(2) {
+                    let start_frame_no = chunk[0];
+                    let end_frame_no = chunk[1];
+                    let timestamp = *self.graph.edges(start_frame_no).find_map(|(_, to, ts)| (to == end_frame_no).then_some(ts)).unwrap();
                     let key = SegmentKey {
-                        start_frame_no: chunk[0],
-                        end_frame_no: chunk[1],
+                        start_frame_no,
+                        end_frame_no,
+                        timestamp,
                     };
                     segments.push(key);
                 }
@@ -448,7 +452,7 @@ fn list_segments<'a>(
 ) -> Result<()> {
     let mut stmt = conn.prepare_cached(
         r#"
-    SELECT created_at, size, start_frame_no, end_frame_no
+    SELECT timestamp, size, start_frame_no, end_frame_no
     FROM segments as s
     JOIN monitored_namespaces as m
     ON m.id == s.namespace_id
@@ -462,9 +466,9 @@ fn list_segments<'a>(
             key: SegmentKey {
                 start_frame_no: r.get(2)?,
                 end_frame_no: r.get(3)?,
+                timestamp: r.get(0)?,
             },
             size: r.get(1)?,
-            created_at: DateTime::from_timestamp(r.get(0)?, 0).unwrap(),
         })
     })?;
 
@@ -486,7 +490,7 @@ fn register_segment_info(
     INSERT OR IGNORE INTO segments (
         start_frame_no,
         end_frame_no,
-        created_at,
+        timestamp,
         size,
         namespace_id
     ) 
@@ -495,7 +499,7 @@ fn register_segment_info(
     stmt.execute((
         info.key.start_frame_no,
         info.key.end_frame_no,
-        info.created_at.timestamp(),
+        info.key.timestamp,
         info.size,
         namespace_id,
     ))?;
@@ -508,7 +512,7 @@ fn segments_range(
 ) -> Result<Option<(SegmentInfo, SegmentInfo)>> {
     let mut stmt = conn.prepare_cached(
         r#"
-    SELECT min(created_at), size, start_frame_no, end_frame_no
+    SELECT min(timestamp), size, start_frame_no, end_frame_no
     FROM segments as s
     JOIN monitored_namespaces as m
     ON m.id == s.namespace_id
@@ -522,16 +526,16 @@ fn segments_range(
                 key: SegmentKey {
                     start_frame_no: r.get(2)?,
                     end_frame_no: r.get(3)?,
+                    timestamp: r.get(0)?,
                 },
                 size: r.get(1)?,
-                created_at: DateTime::from_timestamp(r.get(0)?, 0).unwrap(),
             })
         })
         .optional()?;
 
     let mut stmt = conn.prepare_cached(
         r#"
-    SELECT max(created_at), size, start_frame_no, end_frame_no
+    SELECT max(timestamp), size, start_frame_no, end_frame_no
     FROM segments as s
     JOIN monitored_namespaces as m
     ON m.id == s.namespace_id
@@ -545,9 +549,9 @@ fn segments_range(
                 key: SegmentKey {
                     start_frame_no: r.get(2)?,
                     end_frame_no: r.get(3)?,
+                    timestamp: r.get(0)?,
                 },
                 size: r.get(1)?,
-                created_at: DateTime::from_timestamp(r.get(0)?, 0).unwrap(),
             })
         })
         .optional()?;

libsql-wal/src/storage/mod.rs

@@ -64,6 +64,7 @@ pub enum RestoreOptions {
 pub struct SegmentKey {
     pub start_frame_no: u64,
     pub end_frame_no: u64,
+    pub timestamp: u64,
 }
 
 impl PartialOrd for SegmentKey {
@@ -115,6 +116,7 @@ impl From<&SegmentMeta> for SegmentKey {
         Self {
             start_frame_no: value.start_frame_no,
             end_frame_no: value.end_frame_no,
+            timestamp: value.segment_timestamp.timestamp_millis() as _,
         }
     }
 }
@@ -125,11 +127,13 @@ impl FromStr for SegmentKey {
     fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
         let (rev_start_fno, s) = s.split_at(20);
         let start_frame_no = u64::MAX - rev_start_fno.parse::<u64>().map_err(|_| ())?;
-        let (_, rev_end_fno) = s.split_at(1);
+        let (rev_end_fno, timestamp) = s[1..].split_at(20);
         let end_frame_no = u64::MAX - rev_end_fno.parse::<u64>().map_err(|_| ())?;
+        let timestamp = timestamp[1..].parse().map_err(|_| ())?;
         Ok(Self {
             start_frame_no,
             end_frame_no,
+            timestamp,
         })
     }
 }
@@ -138,9 +142,10 @@ impl fmt::Display for SegmentKey {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "{:020}-{:020}",
+            "{:020}-{:020}-{:020}",
             u64::MAX - self.start_frame_no,
             u64::MAX - self.end_frame_no,
+            self.timestamp,
         )
     }
 }
@@ -219,9 +224,6 @@ pub trait Storage: Send + Sync + 'static {
 pub struct SegmentInfo {
     pub key: SegmentKey,
     pub size: usize,
-    /// when that segment was created. This is different from the segment timestamp, corresponding
-    /// to the last commit date in this segment
-    pub created_at: DateTime<Utc>,
 }
 
 /// special zip function for Either storage implementation
@@ -506,6 +508,7 @@ impl<IO: Io> Storage for TestStorage<IO> {
             let key = SegmentKey {
                 start_frame_no: seg.header().start_frame_no.get(),
                 end_frame_no,
+                timestamp: seg.header().sealed_at_timestamp.get(),
             };
             let index = Map::new(index.into()).unwrap();
             inner