Fix configuration not being synced with replicas (#1071)

Author: MarinPostma

libsql-server/src/namespace/meta_store.rs

@@ -5,6 +5,7 @@ use std::{collections::HashMap, fs::read_dir};
 
 use bottomless::bottomless_wal::BottomlessWalWrapper;
 use bottomless::replicator::CompressionKind;
+use futures_core::Future;
 use libsql_replication::rpc::metadata;
 use libsql_sys::wal::{
     wrapper::{WalWrapper, WrappedWal},
@@ -409,6 +410,18 @@ impl MetaStoreHandle {
         }
     }
 
+    pub fn changed(&self) -> impl Future<Output = ()> {
+        let mut rcv = match &self.inner {
+            HandleState::Internal(_) => panic!("can't wait for change on internal handle"),
+            HandleState::External(_, rcv) => rcv.clone(),
+        };
+        // ack the current value.
+        rcv.borrow_and_update();
+        async move {
+            let _ = rcv.changed().await;
+        }
+    }
+
     pub async fn store(&self, new_config: impl Into<Arc<DatabaseConfig>>) -> Result<()> {
         match &self.inner {
             HandleState::Internal(config) => {
@@ -418,6 +431,8 @@ impl MetaStoreHandle {
                 let new_config = new_config.into();
                 tracing::debug!(?new_config, "storing new namespace config");
                 let mut c = config.clone();
+                // ack the current value.
+                c.borrow_and_update();
                 let changed = c.changed();
 
                 changes_tx

libsql-server/src/namespace/mod.rs

@@ -840,6 +840,10 @@ impl<T: Database> Namespace<T> {
     pub fn stats(&self) -> Arc<Stats> {
         self.stats.clone()
     }
+
+    pub fn config_changed(&self) -> impl Future<Output = ()> {
+        self.db_config_store.changed()
+    }
 }
 
 pub struct ReplicaNamespaceConfig {

libsql-server/src/rpc/replication_log.rs

@@ -6,14 +6,15 @@ use std::sync::{Arc, RwLock};
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use futures::stream::BoxStream;
+use futures_core::Future;
 pub use libsql_replication::rpc::replication as rpc;
 use libsql_replication::rpc::replication::replication_log_server::ReplicationLog;
 use libsql_replication::rpc::replication::{
     Frame, Frames, HelloRequest, HelloResponse, LogOffset, NAMESPACE_DOESNT_EXIST,
     NEED_SNAPSHOT_ERROR_MSG, NO_HELLO_ERROR_MSG, SESSION_TOKEN_KEY,
 };
 use md5::{Digest, Md5};
-use tokio_stream::StreamExt;
+use tokio_stream::StreamExt as _;
 use tonic::transport::server::TcpConnectInfo;
 use tonic::Status;
 use uuid::Uuid;
@@ -139,18 +140,20 @@ impl ReplicationLogService {
             Arc<DatabaseConfig>,
             usize,
             Arc<Stats>,
+            impl Future<Output = ()>,
         ),
         Status,
     > {
-        let (logger, config, version, stats) = self
+        let (logger, config, version, stats, config_changed) = self
             .namespaces
             .with(namespace, |ns| {
                 let logger = ns.db.wal_manager.wrapped().logger().clone();
+                let config_changed = ns.config_changed();
                 let config = ns.config();
                 let version = ns.config_version();
                 let stats = ns.stats();
 
-                (logger, config, version, stats)
+                (logger, config, version, stats, config_changed)
             })
             .await
             .map_err(|e| {
@@ -165,7 +168,7 @@ impl ReplicationLogService {
             self.verify_session_token(req, version)?;
         }
 
-        Ok((logger, config, version, stats))
+        Ok((logger, config, version, stats, config_changed))
     }
 
     fn encode_session_token(&self, version: usize) -> Uuid {
@@ -249,7 +252,8 @@ impl ReplicationLog for ReplicationLogService {
 
         self.authenticate(&req, namespace.clone()).await?;
 
-        let (logger, _, _, stats) = self.logger_from_namespace(namespace, &req, true).await?;
+        let (logger, _, _, stats, config_changed) =
+            self.logger_from_namespace(namespace, &req, true).await?;
 
         let stats = if self.collect_stats {
             Some(stats)
@@ -259,13 +263,29 @@ impl ReplicationLog for ReplicationLogService {
 
         let req = req.into_inner();
 
-        let stream = StreamGuard::new(
+        let mut stream = StreamGuard::new(
             FrameStream::new(logger, req.next_offset, true, None, stats)
                 .map_err(|e| Status::internal(e.to_string()))?,
             self.idle_shutdown_layer.clone(),
         )
         .map(map_frame_stream_output);
 
+        // if only tokio_stream had futures::Stream::take_until...
+        let stream = async_stream::stream! {
+            tokio::pin!(config_changed);
+            loop {
+                tokio::select! {
+                    _ = &mut config_changed => {
+                        break
+                    },
+                    Some(next) = stream.next() => {
+                        yield next
+                    }
+                    else => break,
+                }
+            }
+        };
+
         Ok(tonic::Response::new(Box::pin(stream)))
     }
 
@@ -276,7 +296,7 @@ impl ReplicationLog for ReplicationLogService {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
         self.authenticate(&req, namespace.clone()).await?;
 
-        let (logger, _, _, stats) = self.logger_from_namespace(namespace, &req, true).await?;
+        let (logger, _, _, stats, _) = self.logger_from_namespace(namespace, &req, true).await?;
 
         let stats = if self.collect_stats {
             Some(stats)
@@ -324,7 +344,7 @@ impl ReplicationLog for ReplicationLogService {
             }
         }
 
-        let (logger, config, version, _) =
+        let (logger, config, version, _, _) =
             self.logger_from_namespace(namespace, &req, false).await?;
 
         let session_hash = self.encode_session_token(version);
@@ -348,7 +368,7 @@ impl ReplicationLog for ReplicationLogService {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
         self.authenticate(&req, namespace.clone()).await?;
 
-        let (logger, _, _, stats) = self.logger_from_namespace(namespace, &req, true).await?;
+        let (logger, _, _, stats, _) = self.logger_from_namespace(namespace, &req, true).await?;
 
         let stats = if self.collect_stats {
             Some(stats)