feat(server): improve doc gc (#15363)

#### PR Dependency Tree


* **PR #15363** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Improved validation of workspace roots and document projections, with
clearer failures for malformed or incomplete data.
  * Improved document reference rebuilding and cleanup reliability.
* Updated document update merging to better handle invalid binary data.

* **Performance**
* Avoided unnecessary document reconstruction when no updates are
pending.

* **Tests**
* Updated coverage for malformed workspace roots and document snapshot
parsing.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs

@@ -1,4 +1,5 @@
 use affine_doc_loader as doc_loader;
+use chrono::{DateTime, Utc};
 use sqlx::PgPool;
 
 use super::{
@@ -8,11 +9,7 @@ use super::{
 
 const PARSER_VERSION: i32 = 1;
 
-struct ExtractedRef {
-  blob_key: String,
-  block_id: String,
-  flavour: String,
-}
+type ExtractedRef = doc_loader::BlobRef;
 
 #[derive(Default)]
 struct ProjectionState {
@@ -153,23 +150,9 @@ async fn purge_removed_doc_refs(pool: &PgPool, workspace_id: &str, current_doc_i
   Ok(result.rows_affected() as i64)
 }
 
-fn extract_refs(snapshot: &CurrentDoc) -> RuntimeResult<Vec<ExtractedRef>> {
-  let parsed = doc_loader::parse_doc_from_binary(snapshot.blob.clone(), snapshot.doc_id.clone())
-    .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs parse failed: {err}")))?;
-  let mut refs = Vec::new();
-  for block in parsed.blocks {
-    let Some(blob_keys) = block.blob else {
-      continue;
-    };
-    for blob_key in blob_keys {
-      refs.push(ExtractedRef {
-        blob_key,
-        block_id: block.block_id.clone(),
-        flavour: block.flavour.clone(),
-      });
-    }
-  }
-  Ok(refs)
+fn extract_refs(blob: Vec<u8>) -> RuntimeResult<Vec<ExtractedRef>> {
+  doc_loader::get_blob_refs_from_binary(blob)
+    .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs parse failed: {err}")))
 }
 
 #[cfg(test)]
@@ -191,7 +174,7 @@ mod tests {
       updated_at: Utc::now(),
     };
 
-    let refs = extract_refs(&snapshot).expect("refs should parse");
+    let refs = extract_refs(snapshot.blob).expect("refs should parse");
 
     assert!(
       refs
@@ -233,19 +216,25 @@ mod tests {
       updated_at: Utc::now(),
     };
 
-    assert!(extract_refs(&snapshot).is_err());
+    assert!(extract_refs(snapshot.blob).is_err());
   }
 }
 
-async fn replace_doc_refs(pool: &PgPool, snapshot: &CurrentDoc, refs: Vec<ExtractedRef>) -> RuntimeResult<(i64, i64)> {
+async fn replace_doc_refs(
+  pool: &PgPool,
+  workspace_id: &str,
+  doc_id: &str,
+  updated_at: DateTime<Utc>,
+  refs: Vec<ExtractedRef>,
+) -> RuntimeResult<(i64, i64)> {
   let mut tx = pool
     .begin()
     .await
     .map_err(|err| RuntimeError::database("Doc blob refs transaction failed", err))?;
 
   let deleted = sqlx::query("DELETE FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2")
-    .bind(&snapshot.workspace_id)
-    .bind(&snapshot.doc_id)
+    .bind(workspace_id)
+    .bind(doc_id)
     .execute(&mut *tx)
     .await
     .map_err(|err| RuntimeError::database("Doc blob refs delete failed", err))?
@@ -267,12 +256,12 @@ async fn replace_doc_refs(pool: &PgPool, snapshot: &CurrentDoc, refs: Vec<Extrac
             error = NULL
       "#,
     )
-    .bind(&snapshot.workspace_id)
-    .bind(&snapshot.doc_id)
+    .bind(workspace_id)
+    .bind(doc_id)
     .bind(reference.blob_key)
     .bind(reference.block_id)
     .bind(reference.flavour)
-    .bind(snapshot.updated_at)
+    .bind(updated_at)
     .bind(PARSER_VERSION)
     .execute(&mut *tx)
     .await
@@ -330,9 +319,15 @@ async fn rebuild_doc_blob_refs_inner(
     return Ok(result);
   };
 
-  match extract_refs(&snapshot) {
+  let CurrentDoc {
+    workspace_id,
+    doc_id,
+    blob,
+    updated_at,
+  } = snapshot;
+  match extract_refs(blob) {
     Ok(refs) => {
-      let (written, deleted) = replace_doc_refs(&pool, &snapshot, refs).await?;
+      let (written, deleted) = replace_doc_refs(&pool, &workspace_id, &doc_id, updated_at, refs).await?;
       result.parsed_docs = 1;
       result.refs_written = written;
       result.refs_deleted = deleted;

packages/backend/native/src/runtime/storage_runtime/document_cleanup.rs

@@ -311,9 +311,12 @@ async fn current_activity(
 }
 
 fn root_contains(root: CurrentDoc, doc_id: &str) -> RuntimeResult<bool> {
-  let ids = affine_doc_loader::get_doc_ids_from_binary(root.blob, true)
+  let projection = affine_doc_loader::project_workspace_root(root.blob, true)
     .map_err(|err| RuntimeError::invalid_state(format!("Document cleanup root parse failed: {err}")))?;
-  Ok(ids.iter().any(|id| id == doc_id))
+  if !projection.complete {
+    return Err(RuntimeError::invalid_state("Document cleanup root doc is incomplete"));
+  }
+  Ok(projection.doc_ids.iter().any(|id| id == doc_id))
 }
 
 async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candidate) -> RuntimeResult<i64> {

packages/backend/native/src/runtime/storage_runtime/mod.rs

@@ -101,6 +101,9 @@ fn merge_current_doc(
   if snapshot.is_none() && updates.is_empty() {
     return Ok(None);
   }
+  if updates.is_empty() {
+    return Ok(snapshot);
+  }
   let mut doc = Doc::default();
   let mut updated_at = snapshot
     .as_ref()
@@ -136,8 +139,12 @@ async fn load_workspace_live_doc_ids(pool: &PgPool, workspace_id: &str) -> Runti
 
 fn workspace_live_doc_ids(root: Option<CurrentDoc>) -> RuntimeResult<Vec<String>> {
   let root = root.ok_or_else(|| RuntimeError::invalid_state("Workspace root doc is missing"))?;
-  let mut ids = affine_doc_loader::get_doc_ids_from_binary(root.blob, true)
+  let projection = affine_doc_loader::project_workspace_root(root.blob, true)
     .map_err(|err| RuntimeError::invalid_state(format!("Workspace root doc parse failed: {err}")))?;
+  if !projection.complete {
+    return Err(RuntimeError::invalid_state("Workspace root doc is incomplete"));
+  }
+  let mut ids = projection.doc_ids;
   ids.sort();
   ids.dedup();
   Ok(ids)
@@ -1556,6 +1563,18 @@ mod tests {
       }))
       .is_err()
     );
+    assert!(
+      workspace_live_doc_ids(Some(CurrentDoc {
+        workspace_id: "workspace".to_string(),
+        doc_id: "workspace".to_string(),
+        blob: vec![
+          1, 1, 1, 1, 40, 0, 1, 0, 11, 115, 117, 98, 95, 109, 97, 112, 95, 107, 101, 121, 1, 119, 13, 115, 117, 98, 95,
+          109, 97, 112, 95, 118, 97, 108, 117, 101, 0,
+        ],
+        updated_at: Utc::now(),
+      }))
+      .is_err()
+    );
   }
 
   #[test]

packages/frontend/native/nbstore/src/indexer.rs

@@ -3,7 +3,7 @@ use memory_indexer::{SearchHit, SnapshotData};
 use napi_derive::napi;
 use serde::Serialize;
 use sqlx::Row;
-use y_octo::DocOptions;
+use y_octo::merge_updates_v1;
 
 // Increment this whenever there is a breaking change in the index format or how
 // updates are applied
@@ -120,7 +120,7 @@ impl SqliteDocStorage {
     }
     segments.extend(updates.into_iter().map(|update| update.bin.to_vec()));
 
-    merge_updates(segments, doc_id).map(Some)
+    merge_updates(segments).map(Some)
   }
 
   pub async fn init_index(&self) -> Result<()> {
@@ -249,7 +249,7 @@ impl SqliteDocStorage {
   }
 }
 
-fn merge_updates(mut segments: Vec<Vec<u8>>, guid: &str) -> Result<Vec<u8>> {
+fn merge_updates(mut segments: Vec<Vec<u8>>) -> Result<Vec<u8>> {
   if segments.is_empty() {
     return Err(ParseError::DocNotFound.into());
   }
@@ -258,15 +258,9 @@ fn merge_updates(mut segments: Vec<Vec<u8>>, guid: &str) -> Result<Vec<u8>> {
     return segments.pop().ok_or(ParseError::DocNotFound.into());
   }
 
-  let mut doc = DocOptions::new().with_guid(guid.to_string()).build();
-  for update in segments.iter() {
-    doc
-      .apply_update_from_binary_v1(update)
-      .map_err(|_| ParseError::InvalidBinary)?;
-  }
-
-  let buffer = doc
-    .encode_update_v1()
+  let update = merge_updates_v1(segments).map_err(|_| ParseError::InvalidBinary)?;
+  let buffer = update
+    .encode_v1()
     .map_err(|err| ParseError::ParserError(err.to_string()))?;
 
   Ok(buffer)
@@ -317,6 +311,13 @@ mod tests {
       .execute(&storage.pool)
       .await
       .unwrap();
+    sqlx::query(r#"INSERT INTO updates (doc_id, data, created_at) VALUES (?, ?, ?)"#)
+      .bind("demo-doc")
+      .bind(&[0, 0][..])
+      .bind(Utc::now().naive_utc())
+      .execute(&storage.pool)
+      .await
+      .unwrap();
 
     let result = storage.crawl_doc_data("demo-doc").await.unwrap();