add heads to reads

GET /store?heads=true and GET /store/{key}?heads=true add the header
x-continuum-heads (sorted, comma-separated hex). The heads come from
the same moment as the read.

- It's opt-in because reading the heads takes the write lock.
  Reads without the flag don't change.
- Stale reads during a rejoin never carry heads.

1 files changed, 39 insertions, 5 deletions

diff --git a/src/state.rs b/src/state.rs
index 20285a0..42393b6 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -7,7 +7,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::{bail, Result};
-use automerge::{sync, AutoCommit};
+use automerge::{sync, AutoCommit, ChangeHash};
 use tokio::sync::{broadcast, mpsc, watch, Mutex, RwLock};
 use tokio::time::Instant;
 use tracing::{debug, error, warn};
@@ -573,18 +573,24 @@ 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). That needs the
+    /// write lock, so it is opt-in. 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), axum::http::StatusCode>
+    ) -> Result<(R, ReadSource, Option<Vec<ChangeHash>>), 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 (rendered, heads) = self.render_live(with_heads, read, finish).await;
+            return Ok((rendered, ReadSource::Live, heads));
         }
 
         let stale = self
@@ -595,7 +601,8 @@ impl AppState {
             .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));
+                let (rendered, heads) = self.render_live(with_heads, read, finish).await;
+                return Ok((rendered, ReadSource::Live, heads));
             }
             return Err(axum::http::StatusCode::SERVICE_UNAVAILABLE);
         };
@@ -603,7 +610,34 @@ impl AppState {
         let rendered = tokio::task::spawn_blocking(move || finish(read(&doc)))
             .await
             .expect("document read task panicked");
-        Ok((rendered, ReadSource::Stale(epoch)))
+        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);
+        }
+        // AutoCommit::get_heads takes &mut self (it closes a pending transaction).
+        let (data, mut heads) = self
+            .write_doc(move |doc| {
+                let heads = doc.get_heads();
+                (read(doc), heads)
+            })
+            .await;
+        heads.sort();
+        let rendered = tokio::task::spawn_blocking(move || finish(data))
+            .await
+            .expect("document read task panicked");
+        (rendered, Some(heads))
     }
 
     fn set_stale(&self, view: Option<StaleView>) {