return CompactedSegmentHeader when fetching compacted segment

Author: MarinPostma

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

@@ -8,6 +8,7 @@ use uuid::Uuid;
 
 use super::{RestoreOptions, Result, SegmentKey};
 use crate::io::file::FileExt;
+use crate::segment::compacted::CompactedSegmentDataHeader;
 use libsql_sys::name::NamespaceName;
 
 // pub mod fs;
@@ -63,7 +64,7 @@ pub trait Backend: Send + Sync + 'static {
         namespace: &NamespaceName,
         key: &SegmentKey,
         file: &impl FileExt,
-    ) -> Result<()>;
+    ) -> Result<CompactedSegmentDataHeader>;
 
     // this method taking self: Arc<Self> is an infortunate consequence of rust type system making
     // impl FileExt variant with all the arguments, with no escape hatch...
@@ -176,7 +177,7 @@ impl<T: Backend> Backend for Arc<T> {
         namespace: &NamespaceName,
         key: &SegmentKey,
         file: &impl FileExt,
-    ) -> Result<()> {
+    ) -> Result<CompactedSegmentDataHeader> {
         self.as_ref()
             .fetch_segment_data_to_file(config, namespace, key, file)
             .await

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

@@ -17,7 +17,7 @@ use bytes::{Bytes, BytesMut};
 use http_body::{Frame as HttpFrame, SizeHint};
 use libsql_sys::name::NamespaceName;
 use roaring::RoaringBitmap;
-use tokio::io::{AsyncRead, AsyncReadExt, BufReader};
+use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader};
 use tokio_util::sync::ReusableBoxFuture;
 use zerocopy::byteorder::little_endian::{U16 as lu16, U32 as lu32, U64 as lu64};
 use zerocopy::{AsBytes, FromBytes, FromZeroes};
@@ -135,12 +135,17 @@ impl<IO: Io> S3Backend<IO> {
         folder_key: &FolderKey<'_>,
         segment_key: &SegmentKey,
         file: &impl FileExt,
-    ) -> Result<()> {
+    ) -> Result<CompactedSegmentDataHeader> {
         let reader = self
             .fetch_segment_data_reader(config, folder_key, segment_key)
             .await?;
+        let mut reader = tokio::io::BufReader::with_capacity(8196, reader);
+        while reader.fill_buf().await?.len() < size_of::<CompactedSegmentDataHeader>() {}
+        let header = CompactedSegmentDataHeader::read_from_prefix(reader.buffer()).unwrap();
+
         copy_to_file(reader, file).await?;
-        Ok(())
+
+        Ok(header)
     }
 
     async fn s3_get(&self, config: &S3Config, key: String) -> Result<ByteStream> {
@@ -487,14 +492,15 @@ where
         namespace: &NamespaceName,
         key: &SegmentKey,
         file: &impl FileExt,
-    ) -> Result<()> {
+    ) -> Result<CompactedSegmentDataHeader> {
         let folder_key = FolderKey {
             cluster_id: &config.cluster_id,
             namespace: &namespace,
         };
-        self.fetch_segment_data_inner(config, &folder_key, key, file)
+        let header = self
+            .fetch_segment_data_inner(config, &folder_key, key, file)
             .await?;
-        Ok(())
+        Ok(header)
     }
 
     async fn fetch_segment_data(

libsql-wal/src/storage/job.rs

@@ -116,6 +116,7 @@ mod test {
 
     use crate::io::file::FileExt;
     use crate::io::StdIO;
+    use crate::segment::compacted::CompactedSegmentDataHeader;
     use crate::storage::{RestoreOptions, SegmentKey};
     // use crate::registry::WalRegistry;
     // use crate::segment::compacted::CompactedSegmentDataHeader;
@@ -497,7 +498,7 @@ mod test {
                 _namespace: &NamespaceName,
                 _key: &SegmentKey,
                 _file: &impl FileExt,
-            ) -> Result<()> {
+            ) -> Result<CompactedSegmentDataHeader> {
                 todo!()
             }