split state moduleHEADmaster

1 files changed, 100 insertions, 0 deletions

diff --git a/src/state/read.rs b/src/state/read.rs
new file mode 100644
index 0000000..b76582c
--- /dev/null
+++ b/src/state/read.rs
@@ -0,0 +1,100 @@
+//! Serving reads: from the live document, or from the pre-rejoin copy while rejoining.
+
+use std::sync::Arc;
+
+use automerge::{AutoCommit, ChangeHash};
+
+use super::AppState;
+use crate::document::replica::Epoch;
+
+/// 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`.
+    ///
+    /// With `with_heads`, a live read also returns the document's heads at the moment of
+    /// the read (for clients that base later requests on what they read). Stale reads
+    /// have no heads.
+    pub(crate) async fn render_read<T, R>(
+        &self,
+        unready_live: bool,
+        with_heads: bool,
+        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
+        finish: impl FnOnce(T) -> R + Send + 'static,
+    ) -> Result<(R, ReadSource, Option<Vec<ChangeHash>>), axum::http::StatusCode>
+    where
+        T: Send + 'static,
+        R: Send + 'static,
+    {
+        if self.is_ready().await {
+            let (rendered, heads) = self.render_live(with_heads, read, finish).await;
+            return Ok((rendered, ReadSource::Live, heads));
+        }
+
+        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 {
+                let (rendered, heads) = self.render_live(with_heads, read, finish).await;
+                return Ok((rendered, ReadSource::Live, heads));
+            }
+            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), None))
+    }
+
+    async fn render_live<T, R>(
+        &self,
+        with_heads: bool,
+        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
+        finish: impl FnOnce(T) -> R + Send + 'static,
+    ) -> (R, Option<Vec<ChangeHash>>)
+    where
+        T: Send + 'static,
+        R: Send + 'static,
+    {
+        if !with_heads {
+            return (self.render(read, finish).await, None);
+        }
+        // Heads and data under the same read lock (the heads are cached: `Replica`).
+        let doc = Arc::clone(&self.doc);
+        let (data, heads) = tokio::task::spawn_blocking(move || {
+            let replica = doc.blocking_read();
+            (read(&replica.doc), replica.heads().to_vec())
+        })
+        .await
+        .expect("document read task panicked");
+        let rendered = tokio::task::spawn_blocking(move || finish(data))
+            .await
+            .expect("document read task panicked");
+        (rendered, Some(heads))
+    }
+
+    pub(super) fn set_stale(&self, view: Option<StaleView>) {
+        *self
+            .stale
+            .write()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = view;
+    }
+}