1 files changed, 57 insertions, 3 deletions

diff --git a/src/swarm/document_file.rs b/src/swarm/document_file.rs
index 6a516a1..e9deca0 100644
--- a/src/swarm/document_file.rs
+++ b/src/swarm/document_file.rs
@@ -1,8 +1,15 @@
 //! The worker that keeps the document file (`document::file`) up to date: it looks at
-//! the node (ready, not rejoining, the document's epoch and heads, this node's actors),
-//! and the file says what is due.
+//! the node (ready, the document's epoch and heads, this node's actors), and the file
+//! says what is due. While a ready node rejoins, the file follows the document before
+//! the rejoin, which takes the node's writes meanwhile (`state::StaleView`): a restart
+//! then resumes on it, and rejoins again.
 
+use std::sync::Arc;
+
+use crate::document::Document;
 use crate::document::file::Due;
+use crate::document::replica::Epoch;
+use crate::locks::RwLockExt;
 use crate::state::AppState;
 
 /// Keeps the document file up to date. Runs for as long as the node does.
@@ -23,7 +30,13 @@ pub(crate) async fn flush(state: &AppState) {
     let Some(file) = state.document_file.clone() else {
         return;
     };
-    if !state.is_ready().await || state.remake.is_rejoining() {
+    if state.remake.is_rejoining() {
+        if let Some((epoch, stale)) = state.stale_document() {
+            flush_before_rejoin(&file, epoch, stale).await;
+        }
+        return;
+    }
+    if !state.is_ready().await {
         return;
     }
     let own = state.remake.own_actors();
@@ -54,3 +67,44 @@ pub(crate) async fn flush(state: &AppState) {
         Due::Nothing => {}
     }
 }
+
+/// [`flush`] of the document before a rejoin (`epoch`'s), while it takes this node's
+/// writes. Only appends: a rewrite (another epoch, or the increments outgrew the full
+/// save) waits for the end of the rejoin, so the file keeps the actors it had.
+async fn flush_before_rejoin(
+    file: &Arc<crate::document::file::DocumentFile>,
+    epoch: Epoch,
+    stale: Arc<std::sync::RwLock<Document>>,
+) {
+    let heads = {
+        let stale = Arc::clone(&stale);
+        // Sorted, as the live document's (`Replica::heads`) the file compares with.
+        tokio::task::spawn_blocking(move || {
+            let mut heads = stale.read_or_recover().get_heads();
+            heads.sort_unstable();
+            heads
+        })
+        .await
+        .expect("document read task panicked")
+    };
+    let Some(actors) = file.actors() else {
+        return;
+    };
+    if let Due::Append {
+        epoch: written,
+        since,
+    } = file.due(epoch, actors, &heads)
+        && written == epoch
+    {
+        let since_then = since.clone();
+        let (heads, increment) = tokio::task::spawn_blocking(move || {
+            let doc = stale.read_or_recover();
+            let mut heads = doc.get_heads();
+            heads.sort_unstable();
+            (heads, crate::document::save_since(&doc, &since_then))
+        })
+        .await
+        .expect("document read task panicked");
+        file.append(since, heads, increment).await;
+    }
+}