Make living replica keep primary alive (#517)

Each request for new frames from a replica to a primary
resets the idle timer but such request is a long poll.
It won't return until it gets some new frame.
This means that a replica won't send another request for
new frames in the mean-time and won't restart the idle timer.

This commit makes primary keep track of replicas waiting for
new frames and prevents IdleShutdownLayer from stopping
the primary if there is at least one replica still waiting
for the new frames.

When a replica is put into sleep, it will drop the connection
and the connected_replicas counter will be decreased so only
live replicas can prevent primary from going to sleep.

Fixes #506

Signed-off-by: Piotr Jastrzebski <[email protected]>

Author: haaawk

sqld/src/rpc/mod.rs

@@ -28,7 +28,7 @@ pub async fn run_rpc_server<D: Database>(
     idle_shutdown_layer: Option<IdleShutdownLayer>,
 ) -> anyhow::Result<()> {
     let proxy_service = ProxyService::new(factory, logger.new_frame_notifier.subscribe());
-    let logger_service = ReplicationLogService::new(logger);
+    let logger_service = ReplicationLogService::new(logger, idle_shutdown_layer.clone());
 
     tracing::info!("serving write proxy server at {addr}");
 

sqld/src/rpc/replication_log.rs

@@ -15,23 +15,29 @@ use tonic::Status;
 
 use crate::replication::primary::frame_stream::FrameStream;
 use crate::replication::{LogReadError, ReplicationLogger};
+use crate::utils::services::idle_shutdown::IdleShutdownLayer;
 
 use self::rpc::replication_log_server::ReplicationLog;
 use self::rpc::{Frame, HelloRequest, HelloResponse, LogOffset};
 
 pub struct ReplicationLogService {
     logger: Arc<ReplicationLogger>,
     replicas_with_hello: RwLock<HashSet<SocketAddr>>,
+    idle_shutdown_layer: Option<IdleShutdownLayer>,
 }
 
 pub const NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
 pub const NEED_SNAPSHOT_ERROR_MSG: &str = "NEED_SNAPSHOT";
 
 impl ReplicationLogService {
-    pub fn new(logger: Arc<ReplicationLogger>) -> Self {
+    pub fn new(
+        logger: Arc<ReplicationLogger>,
+        idle_shutdown_layer: Option<IdleShutdownLayer>,
+    ) -> Self {
         Self {
             logger,
             replicas_with_hello: RwLock::new(HashSet::<SocketAddr>::new()),
+            idle_shutdown_layer,
         }
     }
 }
@@ -56,6 +62,42 @@ fn map_frame_stream_output(
     }
 }
 
+pub struct StreamGuard<S> {
+    s: S,
+    idle_shutdown_layer: Option<IdleShutdownLayer>,
+}
+
+impl<S> StreamGuard<S> {
+    fn new(s: S, mut idle_shutdown_layer: Option<IdleShutdownLayer>) -> Self {
+        if let Some(isl) = idle_shutdown_layer.as_mut() {
+            isl.add_connected_replica()
+        }
+        Self {
+            s,
+            idle_shutdown_layer,
+        }
+    }
+}
+
+impl<S> Drop for StreamGuard<S> {
+    fn drop(&mut self) {
+        if let Some(isl) = self.idle_shutdown_layer.as_mut() {
+            isl.remove_connected_replica()
+        }
+    }
+}
+
+impl<S: futures::stream::Stream + Unpin> futures::stream::Stream for StreamGuard<S> {
+    type Item = S::Item;
+
+    fn poll_next(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.get_mut().s.poll_next_unpin(cx)
+    }
+}
+
 #[tonic::async_trait]
 impl ReplicationLog for ReplicationLogService {
     type LogEntriesStream = BoxStream<'static, Result<Frame, Status>>;
@@ -75,9 +117,12 @@ impl ReplicationLog for ReplicationLogService {
             }
         }
 
-        let stream = FrameStream::new(self.logger.clone(), req.into_inner().next_offset)
-            .map(map_frame_stream_output)
-            .boxed();
+        let stream = StreamGuard::new(
+            FrameStream::new(self.logger.clone(), req.into_inner().next_offset),
+            self.idle_shutdown_layer.clone(),
+        )
+        .map(map_frame_stream_output)
+        .boxed();
 
         Ok(tonic::Response::new(stream))
     }

sqld/src/utils/services/idle_shutdown.rs

@@ -1,3 +1,4 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -9,11 +10,14 @@ use tower::{Layer, Service};
 #[derive(Clone)]
 pub struct IdleShutdownLayer {
     watcher: Arc<watch::Sender<()>>,
+    connected_replicas: Arc<AtomicUsize>,
 }
 
 impl IdleShutdownLayer {
     pub fn new(idle_timeout: Duration, shutdown_notifier: mpsc::Sender<()>) -> Self {
         let (sender, mut receiver) = watch::channel(());
+        let connected_replicas = Arc::new(AtomicUsize::new(0));
+        let connected_replicas_clone = connected_replicas.clone();
         tokio::spawn(async move {
             loop {
                 // FIXME: if we measure that this is causing performance issues, we may want to
@@ -23,6 +27,9 @@ impl IdleShutdownLayer {
                     Ok(Ok(_)) => continue,
                     Ok(Err(_)) => break,
                     Err(_) => {
+                        if connected_replicas_clone.load(Ordering::SeqCst) > 0 {
+                            continue;
+                        }
                         tracing::info!(
                             "Idle timeout, no new connection in {idle_timeout:.0?}. Shutting down.",
                         );
@@ -39,9 +46,18 @@ impl IdleShutdownLayer {
 
         Self {
             watcher: Arc::new(sender),
+            connected_replicas,
         }
     }
 
+    pub fn add_connected_replica(&mut self) {
+        self.connected_replicas.fetch_add(1, Ordering::SeqCst);
+    }
+
+    pub fn remove_connected_replica(&mut self) {
+        self.connected_replicas.fetch_sub(1, Ordering::SeqCst);
+    }
+
     pub fn into_kicker(self) -> IdleKicker {
         IdleKicker {
             sender: self.watcher,