add ability to list segments

Author: MarinPostma

libsql-wal/src/storage/async_storage.rs

@@ -7,6 +7,7 @@ use chrono::Utc;
 use libsql_sys::name::NamespaceName;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinSet;
+use tokio_stream::Stream;
 
 use crate::io::{FileExt, Io, StdIO};
 use crate::segment::compacted::CompactedSegment;
@@ -267,6 +268,16 @@ where
         let segment = CompactedSegment::open(file).await?;
         Ok(segment)
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        namespace: &'a NamespaceName,
+        until: u64,
+        config_override: Option<Self::Config>,
+    ) -> impl Stream<Item = super::Result<super::SegmentInfo>> + 'a {
+        let config = config_override.unwrap_or_else(|| self.backend.default_config());
+        self.backend.list_segments(config, namespace, until)
+    }
 }
 
 pub struct AsyncStorageInitConfig<B> {

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

@@ -4,9 +4,10 @@ use std::{future::Future, path::Path};
 
 use chrono::{DateTime, Utc};
 use fst::Map;
+use tokio_stream::Stream;
 use uuid::Uuid;
 
-use super::{RestoreOptions, Result, SegmentKey};
+use super::{RestoreOptions, Result, SegmentInfo, SegmentKey};
 use crate::io::file::FileExt;
 use crate::segment::compacted::CompactedSegmentDataHeader;
 use libsql_sys::name::NamespaceName;
@@ -99,6 +100,13 @@ pub trait Backend: Send + Sync + 'static {
         dest: impl FileExt,
     ) -> Result<()>;
 
+    fn list_segments<'a>(
+        &'a self,
+        config: Self::Config,
+        namespace: &'a NamespaceName,
+        until: u64,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a;
+
     /// Returns the default configuration for this storage
     fn default_config(&self) -> Self::Config;
 }
@@ -195,4 +203,13 @@ impl<T: Backend> Backend for Arc<T> {
             .fetch_segment_data(config, namespace, key)
             .await
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        config: Self::Config,
+        namespace: &'a NamespaceName,
+        until: u64,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        self.as_ref().list_segments(config, namespace, until)
+    }
 }

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

@@ -13,11 +13,13 @@ use aws_sdk_s3::operation::create_bucket::CreateBucketError;
 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 http_body::{Frame as HttpFrame, SizeHint};
 use libsql_sys::name::NamespaceName;
 use roaring::RoaringBitmap;
 use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader};
+use tokio_stream::Stream;
 use tokio_util::sync::ReusableBoxFuture;
 use zerocopy::byteorder::little_endian::{U16 as lu16, U32 as lu32, U64 as lu64};
 use zerocopy::{AsBytes, FromBytes, FromZeroes};
@@ -28,7 +30,7 @@ use crate::io::compat::copy_to_file;
 use crate::io::{FileExt, Io, StdIO};
 use crate::segment::compacted::CompactedSegmentDataHeader;
 use crate::segment::Frame;
-use crate::storage::{Error, RestoreOptions, Result, SegmentKey};
+use crate::storage::{Error, RestoreOptions, Result, SegmentInfo, SegmentKey};
 use crate::LIBSQL_MAGIC;
 
 pub struct S3Backend<IO> {
@@ -306,6 +308,52 @@ impl<IO: Io> S3Backend<IO> {
 
         Ok(index)
     }
+
+    fn list_segments_inner<'a>(
+        &'a self,
+        config: Arc<S3Config>,
+        namespace: &'a NamespaceName,
+        _until: u64,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        async_stream::try_stream! {
+            let folder_key = FolderKey { cluster_id: &config.cluster_id, namespace };
+            let lookup_key_prefix = s3_segment_index_lookup_key_prefix(&folder_key);
+
+            let mut continuation_token = None;
+            loop {
+                let objects = self
+                    .client
+                    .list_objects_v2()
+                    .bucket(&config.bucket)
+                    .prefix(lookup_key_prefix.clone())
+                    .set_continuation_token(continuation_token.take())
+                    .send()
+                    .await
+                    .map_err(|e| Error::unhandled(e, "failed to list bucket"))?;
+
+                for entry in objects.contents() {
+                    let key = entry.key().expect("misssing key?");
+                    let key_path: &Path = key.as_ref();
+                    let Some(key) = SegmentKey::validate_from_path(key_path, &folder_key.namespace) else { continue };
+
+                    let infos = SegmentInfo {
+                        key,
+                        size: entry.size().unwrap_or(0) as usize,
+                        created_at: entry.last_modified().unwrap().to_chrono_utc().unwrap(),
+                    };
+
+                    yield infos;
+                }
+
+                if objects.is_truncated().unwrap_or(false) {
+                    assert!(objects.next_continuation_token.is_some());
+                    continuation_token = objects.next_continuation_token;
+                } else {
+                    break
+                }
+            }
+        }
+    }
 }
 
 pub struct S3Config {
@@ -514,6 +562,15 @@ where
             .await?;
         Ok(file)
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        config: Self::Config,
+        namespace: &'a NamespaceName,
+        until: u64,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        self.list_segments_inner(config, namespace, until)
+    }
 }
 
 #[derive(Clone, Copy)]

libsql-wal/src/storage/job.rs

@@ -520,6 +520,16 @@ mod test {
             ) -> Result<fst::Map<Arc<[u8]>>> {
                 todo!()
             }
+
+            fn list_segments<'a>(
+                &'a self,
+                _config: Self::Config,
+                _namespace: &'a NamespaceName,
+                _until: u64,
+            ) -> impl tokio_stream::Stream<Item = Result<crate::storage::SegmentInfo>> + 'a
+            {
+                tokio_stream::iter(std::iter::from_fn(|| todo!()))
+            }
         }
 
         let job = Job {

libsql-wal/src/storage/mod.rs

@@ -1,17 +1,17 @@
 use std::collections::BTreeMap;
-use std::fmt;
-use std::future::Future;
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::Arc;
+use std::{fmt, future::Future};
 
 use chrono::{DateTime, Utc};
 use fst::Map;
 use hashbrown::HashMap;
 use libsql_sys::name::NamespaceName;
 use libsql_sys::wal::either::Either;
 use tempfile::{tempdir, TempDir};
+use tokio_stream::Stream;
 
 use crate::io::{FileExt, Io, StdIO};
 use crate::segment::compacted::CompactedSegment;
@@ -205,6 +205,20 @@ pub trait Storage: Send + Sync + 'static {
     fn shutdown(&self) -> impl Future<Output = ()> + Send {
         async { () }
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        namespace: &'a NamespaceName,
+        until: u64,
+        config_override: Option<Self::Config>,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a;
+}
+
+#[derive(Debug)]
+pub struct SegmentInfo {
+    pub key: SegmentKey,
+    pub size: usize,
+    pub created_at: DateTime<Utc>,
 }
 
 /// special zip function for Either storage implementation
@@ -323,6 +337,22 @@ where
             Either::B(b) => b.shutdown().await,
         }
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        namespace: &'a NamespaceName,
+        until: u64,
+        config_override: Option<Self::Config>,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        match zip(self, config_override) {
+            Either::A((s, c)) => {
+                tokio_util::either::Either::Left(s.list_segments(namespace, until, c))
+            }
+            Either::B((s, c)) => {
+                tokio_util::either::Either::Right(s.list_segments(namespace, until, c))
+            }
+        }
+    }
 }
 
 /// a placeholder storage that doesn't store segment
@@ -388,6 +418,17 @@ impl Storage for NoStorage {
         #[allow(unreachable_code)]
         Result::<CompactedSegment<std::fs::File>>::Err(Error::InvalidIndex(""))
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        _namespace: &'a NamespaceName,
+        _until: u64,
+        _config_override: Option<Self::Config>,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        unimplemented!("no storage!");
+        #[allow(unreachable_code)]
+        tokio_stream::empty()
+    }
 }
 
 #[doc(hidden)]
@@ -555,6 +596,17 @@ impl<IO: Io> Storage for TestStorage<IO> {
             panic!("not storing")
         }
     }
+
+    fn list_segments<'a>(
+        &'a self,
+        _namespace: &'a NamespaceName,
+        _until: u64,
+        _config_override: Option<Self::Config>,
+    ) -> impl Stream<Item = Result<SegmentInfo>> + 'a {
+        todo!();
+        #[allow(unreachable_code)]
+        tokio_stream::empty()
+    }
 }
 
 pub struct StoreSegmentRequest<S, C> {