Merge pull request #1616 from tursodatabase/type-erase-override-config

type erase override config

Author: MarinPostma

libsql-wal/src/storage/async_storage.rs

@@ -1,6 +1,7 @@
 //! `AsyncStorage` is a `Storage` implementation that defer storage to a background thread. The
 //! durable frame_no is notified asynchronously.
 
+use std::any::Any;
 use std::sync::Arc;
 
 use chrono::Utc;
@@ -22,9 +23,9 @@ use super::{OnStoreCallback, RestoreOptions, Storage, StoreSegmentRequest};
 ///
 /// On shutdown, attempts to empty the queue, and flush the receiver. When the last handle of the
 /// receiver is dropped, and the queue is empty, exit.
-pub struct AsyncStorageLoop<B: Backend, IO: Io, S> {
-    receiver: mpsc::UnboundedReceiver<StorageLoopMessage<B::Config, S>>,
-    scheduler: Scheduler<B::Config, S>,
+pub struct AsyncStorageLoop<B, IO: Io, S> {
+    receiver: mpsc::UnboundedReceiver<StorageLoopMessage<S>>,
+    scheduler: Scheduler<S>,
     backend: Arc<B>,
     io: Arc<IO>,
     max_in_flight: usize,
@@ -112,14 +113,28 @@ where
     fn fetch_durable_frame_no_async(
         &self,
         namespace: NamespaceName,
-        ret: oneshot::Sender<u64>,
-        config_override: Option<Arc<B::Config>>,
+        ret: oneshot::Sender<super::Result<u64>>,
+        config_override: Option<Arc<dyn Any + Send + Sync>>,
     ) {
         let backend = self.backend.clone();
-        let config = config_override.unwrap_or_else(|| backend.default_config());
+        let config = match config_override
+            .map(|c| c.downcast::<B::Config>())
+            .transpose()
+        {
+            Ok(Some(config)) => config,
+            Ok(None) => backend.default_config(),
+            Err(_) => {
+                let _ = ret.send(Err(super::Error::InvalidConfigType));
+                return;
+            }
+        };
+
         tokio::spawn(async move {
-            let meta = backend.meta(&config, &namespace).await.unwrap();
-            let _ = ret.send(meta.max_frame_no);
+            let res = backend
+                .meta(&config, &namespace)
+                .await
+                .map(|meta| meta.max_frame_no);
+            let _ = ret.send(res);
         });
     }
 }
@@ -132,18 +147,18 @@ pub struct BottomlessConfig<C> {
     pub config: C,
 }
 
-enum StorageLoopMessage<C, S> {
-    StoreReq(StoreSegmentRequest<C, S>),
+enum StorageLoopMessage<S> {
+    StoreReq(StoreSegmentRequest<S>),
     DurableFrameNoReq {
         namespace: NamespaceName,
-        config_override: Option<Arc<C>>,
-        ret: oneshot::Sender<u64>,
+        config_override: Option<Arc<dyn Any + Send + Sync>>,
+        ret: oneshot::Sender<super::Result<u64>>,
     },
 }
 
-pub struct AsyncStorage<B: Backend, S> {
+pub struct AsyncStorage<B, S> {
     /// send request to the main loop
-    job_sender: mpsc::UnboundedSender<StorageLoopMessage<B::Config, S>>,
+    job_sender: mpsc::UnboundedSender<StorageLoopMessage<S>>,
     force_shutdown: oneshot::Sender<()>,
     backend: Arc<B>,
 }
@@ -163,11 +178,15 @@ where
         config_override: Option<Arc<Self::Config>>,
         on_store_callback: OnStoreCallback,
     ) {
+        fn into_any<T: Sync + Send + 'static>(t: Arc<T>) -> Arc<dyn Any + Sync + Send> {
+            t
+        }
+
         let req = StoreSegmentRequest {
             namespace: namespace.clone(),
             segment,
             created_at: Utc::now(),
-            storage_config_override: config_override,
+            storage_config_override: config_override.map(into_any),
             on_store_callback,
         };
 

libsql-wal/src/storage/error.rs

@@ -19,6 +19,8 @@ pub enum Error {
     // We may recover from this error, and rebuild the index from the data file.
     #[error("invalid index: {0}")]
     InvalidIndex(&'static str),
+    #[error("Provided config is of an invalid type")]
+    InvalidConfigType,
 }
 
 impl Error {

libsql-wal/src/storage/job.rs

@@ -9,13 +9,13 @@ use crate::segment::Segment;
 
 /// A request, with an id
 #[derive(Debug)]
-pub(crate) struct IndexedRequest<C, T> {
-    pub(crate) request: StoreSegmentRequest<C, T>,
+pub(crate) struct IndexedRequest<T> {
+    pub(crate) request: StoreSegmentRequest<T>,
     pub(crate) id: u64,
 }
 
-impl<C, T> Deref for IndexedRequest<C, T> {
-    type Target = StoreSegmentRequest<C, T>;
+impl<T> Deref for IndexedRequest<T> {
+    type Target = StoreSegmentRequest<T>;
 
     fn deref(&self) -> &Self::Target {
         &self.request
@@ -24,10 +24,10 @@ impl<C, T> Deref for IndexedRequest<C, T> {
 
 /// A storage Job to be performed
 #[derive(Debug)]
-pub(crate) struct Job<C, T> {
+pub(crate) struct Job<T> {
     /// Segment to store.
     // TODO: implement request batching (merge segment and send).
-    pub(crate) request: IndexedRequest<C, T>,
+    pub(crate) request: IndexedRequest<T>,
 }
 
 // #[repr(transparent)]
@@ -42,14 +42,14 @@ pub(crate) struct Job<C, T> {
 //     }
 // }
 //
-impl<C, Seg> Job<C, Seg>
+impl<Seg> Job<Seg>
 where
     Seg: Segment,
 {
     /// Perform the job and return the JobResult. This is not allowed to panic.
-    pub(crate) async fn perform<B, IO>(self, backend: B, io: IO) -> JobResult<C, Seg>
+    pub(crate) async fn perform<B, IO>(self, backend: B, io: IO) -> JobResult<Seg>
     where
-        B: Backend<Config = C>,
+        B: Backend,
         IO: Io,
     {
         let result = self.try_perform(backend, io).await;
@@ -58,7 +58,7 @@ where
 
     async fn try_perform<B, IO>(&self, backend: B, io: IO) -> Result<u64>
     where
-        B: Backend<Config = C>,
+        B: Backend,
         IO: Io,
     {
         let segment = &self.request.segment;
@@ -81,6 +81,9 @@ where
             .request
             .storage_config_override
             .clone()
+            .map(|c| c.downcast::<B::Config>())
+            .transpose()
+            .map_err(|_| super::Error::InvalidConfigType)?
             .unwrap_or_else(|| backend.default_config());
 
         backend.store(&config, meta, tmp, new_index).await?;
@@ -90,9 +93,9 @@ where
 }
 
 #[derive(Debug)]
-pub(crate) struct JobResult<C, S> {
+pub(crate) struct JobResult<S> {
     /// The job that was performed
-    pub(crate) job: Job<C, S>,
+    pub(crate) job: Job<S>,
     /// The outcome of the job: the new durable index, or an error.
     pub(crate) result: Result<u64>,
 }

libsql-wal/src/storage/mod.rs

@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::collections::BTreeMap;
 use std::fmt;
 use std::future::Future;
@@ -437,7 +438,7 @@ impl<IO: Io> Storage for TestStorage<IO> {
     }
 }
 
-pub struct StoreSegmentRequest<C, S> {
+pub struct StoreSegmentRequest<S> {
     namespace: NamespaceName,
     /// Path to the segment. Read-only for bottomless
     segment: S,
@@ -446,22 +447,20 @@ pub struct StoreSegmentRequest<C, S> {
 
     /// alternative configuration to use with the storage layer.
     /// e.g: S3 overrides
-    storage_config_override: Option<Arc<C>>,
+    storage_config_override: Option<Arc<dyn Any + Send + Sync>>,
     /// Called after the segment was stored, with the new durable index
     on_store_callback: OnStoreCallback,
 }
 
-impl<C, S> fmt::Debug for StoreSegmentRequest<C, S>
+impl<S> fmt::Debug for StoreSegmentRequest<S>
 where
-    C: fmt::Debug,
     S: fmt::Debug,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("StoreSegmentRequest")
             .field("namespace", &self.namespace)
             .field("segment", &self.segment)
             .field("created_at", &self.created_at)
-            .field("storage_config_override", &self.storage_config_override)
             .finish()
     }
 }

libsql-wal/src/storage/scheduler.rs

@@ -5,13 +5,13 @@ use super::job::{IndexedRequest, Job, JobResult};
 use super::StoreSegmentRequest;
 use libsql_sys::name::NamespaceName;
 
-struct NamespaceRequests<C, F> {
-    requests: VecDeque<IndexedRequest<C, F>>,
+struct NamespaceRequests<F> {
+    requests: VecDeque<IndexedRequest<F>>,
     /// there's work in flight for this namespace
     in_flight: bool,
 }
 
-impl<C, F> Default for NamespaceRequests<C, F> {
+impl<F> Default for NamespaceRequests<F> {
     fn default() -> Self {
         Self {
             requests: Default::default(),
@@ -28,14 +28,14 @@ impl<C, F> Default for NamespaceRequests<C, F> {
 /// processed, because only the most recent segment is checked for durability. This property
 /// ensures that all segments are present up to the max durable index.
 /// It is generic over C: the storage config type (for config overrides), and T, the segment type
-pub(crate) struct Scheduler<C, T> {
+pub(crate) struct Scheduler<T> {
     /// notify new durability index for namespace
-    requests: HashMap<NamespaceName, NamespaceRequests<C, T>>,
+    requests: HashMap<NamespaceName, NamespaceRequests<T>>,
     queue: priority_queue::PriorityQueue<NamespaceName, Reverse<u64>>,
     next_request_id: u64,
 }
 
-impl<C, T> Scheduler<C, T> {
+impl<T> Scheduler<T> {
     pub fn new() -> Self {
         Self {
             requests: Default::default(),
@@ -46,7 +46,7 @@ impl<C, T> Scheduler<C, T> {
 
     /// Register a new request with the scheduler
     #[tracing::instrument(skip_all)]
-    pub fn register(&mut self, request: StoreSegmentRequest<C, T>) {
+    pub fn register(&mut self, request: StoreSegmentRequest<T>) {
         // invariant: new segment comes immediately after the latest segment for that namespace. This means:
         // - immediately after the last registered segment, if there is any
         // - immediately after the last durable index
@@ -71,7 +71,7 @@ impl<C, T> Scheduler<C, T> {
     /// be scheduled, and returns description of the job to be performed. No other job for this
     /// namespace will be scheduled, until the `JobResult` is reported
     #[tracing::instrument(skip_all)]
-    pub fn schedule(&mut self) -> Option<Job<C, T>> {
+    pub fn schedule(&mut self) -> Option<Job<T>> {
         let (name, _) = self.queue.pop()?;
         let requests = self
             .requests
@@ -90,7 +90,7 @@ impl<C, T> Scheduler<C, T> {
     /// Report the job result to the scheduler. If the job result was a success, the request as
     /// removed from the queue, else, the job is rescheduled
     #[tracing::instrument(skip_all, fields(req_id = result.job.request.id))]
-    pub async fn report(&mut self, result: JobResult<C, T>) {
+    pub async fn report(&mut self, result: JobResult<T>) {
         // re-schedule, or report new max durable frame_no for segment
         let name = result.job.request.request.namespace.clone();
         let requests = self
@@ -151,7 +151,7 @@ mod test {
 
     #[tokio::test]
     async fn schedule_simple() {
-        let mut scheduler = Scheduler::<(), ()>::new();
+        let mut scheduler = Scheduler::<()>::new();
 
         let ns1 = NamespaceName::from("test1");
         let ns2 = NamespaceName::from("test2");
@@ -224,7 +224,7 @@ mod test {
 
     #[tokio::test]
     async fn job_error_reschedule() {
-        let mut scheduler = Scheduler::<(), ()>::new();
+        let mut scheduler = Scheduler::<()>::new();
 
         let ns1 = NamespaceName::from("test1");
         let ns2 = NamespaceName::from("test2");
@@ -264,7 +264,7 @@ mod test {
 
     #[tokio::test]
     async fn schedule_while_in_flight() {
-        let mut scheduler = Scheduler::<(), ()>::new();
+        let mut scheduler = Scheduler::<()>::new();
 
         let ns1 = NamespaceName::from("test1");