1 files changed, 24 insertions, 10 deletions

diff --git a/src/state/read.rs b/src/state/read.rs
index f8323a2..e325e61 100644
--- a/src/state/read.rs
+++ b/src/state/read.rs
@@ -1,17 +1,20 @@
-//! Serving reads: from the live document, or from the pre-rejoin copy while rejoining.
+//! Serving reads: from the live document, or from the pre-rejoin document while
+//! rejoining.
 
 use std::sync::Arc;
 
-use automerge::{AutoCommit, ChangeHash};
-
 use super::AppState;
 use crate::document::replica::Epoch;
+use crate::document::{ChangeHash, Document};
 use crate::locks::RwLockExt;
 
-/// A read-only view of the document the node served before it started rejoining.
+/// The document the node served before it started rejoining, of the epoch it left:
+/// reads are answered from it (marked stale), and writes go into it
+/// (`AppState::mutate_doc`), until the rejoin replays them into the new epoch or
+/// restores this document (`Host::finish_rejoin`, `Host::abandon_rejoin`).
 pub(crate) struct StaleView {
     pub(crate) epoch: Epoch,
-    pub(crate) doc: Arc<AutoCommit>,
+    pub(crate) doc: Arc<std::sync::RwLock<Document>>,
 }
 
 /// Where a read was served from.
@@ -33,7 +36,7 @@ impl AppState {
         &self,
         unready_live: bool,
         with_heads: bool,
-        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
+        read: impl FnOnce(&Document) -> T + Send + 'static,
         finish: impl FnOnce(T) -> R + Send + 'static,
     ) -> Result<(R, ReadSource, Option<Vec<ChangeHash>>), axum::http::StatusCode>
     where
@@ -58,16 +61,19 @@ impl AppState {
             return Err(axum::http::StatusCode::SERVICE_UNAVAILABLE);
         };
 
-        let rendered = tokio::task::spawn_blocking(move || finish(read(&doc)))
-            .await
-            .expect("document read task panicked");
+        let rendered = tokio::task::spawn_blocking(move || {
+            let data = read(&doc.read_or_recover());
+            finish(data)
+        })
+        .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,
+        read: impl FnOnce(&Document) -> T + Send + 'static,
         finish: impl FnOnce(T) -> R + Send + 'static,
     ) -> (R, Option<Vec<ChangeHash>>)
     where
@@ -94,4 +100,12 @@ impl AppState {
     pub(super) fn set_stale(&self, view: Option<StaleView>) {
         *self.stale.write_or_recover() = view;
     }
+
+    /// The stale view's document and epoch, if there is one.
+    pub(crate) fn stale_document(&self) -> Option<(Epoch, Arc<std::sync::RwLock<Document>>)> {
+        self.stale
+            .read_or_recover()
+            .as_ref()
+            .map(|view| (view.epoch, Arc::clone(&view.doc)))
+    }
 }