serve stale reads during rejoinHEADmaster

1 files changed, 69 insertions, 0 deletions

diff --git a/src/state.rs b/src/state.rs
index a8b573b..e7c1767 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -300,6 +300,10 @@ pub(crate) struct AppState {
 
     /// Per-node compaction state (prebuilt snapshots, switch status, archive).
     pub(crate) compaction: Arc<Compaction>,
+
+    /// While rejoining a newer epoch, the previous document, served read-only (and
+    /// marked stale) instead of answering 503.
+    pub(crate) stale: Arc<std::sync::RwLock<Option<StaleView>>>,
 }
 
 impl AppState {
@@ -552,6 +556,64 @@ impl AppState {
     }
 }
 
+/// A read-only view of the document the node served before it started rejoining.
+pub(crate) struct StaleView {
+    pub(crate) epoch: Epoch,
+    pub(crate) doc: Arc<AutoCommit>,
+}
+
+/// Where a read was served from.
+pub(crate) enum ReadSource {
+    Live,
+    /// The pre-rejoin document of this epoch.
+    Stale(Epoch),
+}
+
+impl AppState {
+    /// Renders a read for the store API: from the live document when the node is ready,
+    /// from the stale view while it rejoins, and otherwise `503`, or the live document
+    /// anyway with `unready_live`.
+    pub(crate) async fn render_read<T, R>(
+        &self,
+        unready_live: bool,
+        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
+        finish: impl FnOnce(T) -> R + Send + 'static,
+    ) -> Result<(R, ReadSource), axum::http::StatusCode>
+    where
+        T: Send + 'static,
+        R: Send + 'static,
+    {
+        if self.is_ready().await {
+            return Ok((self.render(read, finish).await, ReadSource::Live));
+        }
+
+        let stale = self
+            .stale
+            .read()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .as_ref()
+            .map(|view| (view.epoch, Arc::clone(&view.doc)));
+        let Some((epoch, doc)) = stale else {
+            if unready_live {
+                return Ok((self.render(read, finish).await, ReadSource::Live));
+            }
+            return Err(axum::http::StatusCode::SERVICE_UNAVAILABLE);
+        };
+
+        let rendered = tokio::task::spawn_blocking(move || finish(read(&doc)))
+            .await
+            .expect("document read task panicked");
+        Ok((rendered, ReadSource::Stale(epoch)))
+    }
+
+    fn set_stale(&self, view: Option<StaleView>) {
+        *self
+            .stale
+            .write()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = view;
+    }
+}
+
 pub(crate) fn unix_millis() -> u64 {
     std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
@@ -638,6 +700,13 @@ impl Host for AppState {
         }
     }
 
+    fn set_stale_view(&self, view: Option<(Epoch, AutoCommit)>) {
+        self.set_stale(view.map(|(epoch, doc)| StaleView {
+            epoch,
+            doc: Arc::new(doc),
+        }));
+    }
+
     fn initial_sync_done(&self) -> bool {
         self.initial_sync_done
             .load(std::sync::atomic::Ordering::Acquire)