chore: add collab version into persistence

Author: Horusiath

collab-database/tests/user_test/helper.rs

@@ -76,14 +76,15 @@ impl DatabaseCollabPersistenceService for TestUserDatabasePersistenceImpl {
 
   fn upsert_collab(
     &self,
-    object_id: &collab_entity::uuid_validation::ObjectId,
+    object_id: &ObjectId,
     encoded_collab: EncodedCollab,
   ) -> Result<(), DatabaseError> {
     let db_write = self.db.write_txn();
     let _ = db_write.upsert_doc_with_doc_state(
       self.uid,
       &self.workspace_id,
       &object_id.to_string(),
+      encoded_collab.collab_version.as_ref(),
       encoded_collab.state_vector.to_vec(),
       encoded_collab.doc_state.to_vec(),
     );

collab-plugins/src/cloud_storage/remote_collab.rs

@@ -6,7 +6,9 @@ use std::time::{Duration, SystemTime};
 
 use anyhow::{Error, anyhow};
 use async_trait::async_trait;
-use collab::core::collab::{CollabOptions, DataSource, TransactionMutExt, default_client_id};
+use collab::core::collab::{
+  CollabOptions, DataSource, TransactionMutExt, VersionedData, default_client_id,
+};
 use collab::core::collab_state::SyncState;
 use collab::core::origin::CollabOrigin;
 use collab::lock::RwLock;
@@ -210,8 +212,8 @@ impl RemoteCollab {
   /// If the remote collab contains any updates, it will return None.
   /// Otherwise, it will merge the updates into one and return the merged update.
   #[allow(dead_code)]
-  pub async fn sync(&self, local_collab: Weak<RwLock<Collab>>) -> Result<Vec<u8>, Error> {
-    let mut remote_update = vec![];
+  pub async fn sync(&self, local_collab: Weak<RwLock<Collab>>) -> Result<VersionedData, Error> {
+    let mut remote_update = None;
     // It would be better if creating a edge function that calculate the diff between the local and remote.
     // The local only need to send its state vector to the remote. In this way, the local does not need to
     // get all the updates from remote.
@@ -232,7 +234,7 @@ impl RemoteCollab {
           } else {
             tracing::error!("🔴decode update failed");
           }
-          remote_update = doc_state;
+          remote_update = Some(doc_state);
         },
         DataSource::DocStateV2(doc_state) => {
           if let Ok(update) = Update::decode_v2(&doc_state) {
@@ -242,7 +244,7 @@ impl RemoteCollab {
           } else {
             tracing::error!("🔴decode update failed");
           }
-          remote_update = doc_state;
+          remote_update = Some(doc_state);
         },
       }
 
@@ -311,7 +313,7 @@ impl RemoteCollab {
         object_id_string: self.object.object_id.to_string(),
       });
     }
-    Ok(remote_update)
+    Ok(remote_update.unwrap_or(VersionedData::new(vec![], None)))
   }
 
   pub fn push_update(&self, update: &[u8]) -> Result<(), Error> {

collab-plugins/src/local_storage/indexeddb/indexeddb_plugin.rs

@@ -8,7 +8,7 @@ use collab_entity::CollabType;
 use futures::stream::StreamExt;
 
 use crate::local_storage::kv::PersistenceError;
-use collab::core::collab::make_yrs_doc;
+use collab::core::collab::{CollabContext, make_yrs_doc};
 use collab::core::transaction::DocTransactionExtension;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering::SeqCst;
@@ -64,15 +64,17 @@ impl IndexeddbDiskPlugin {
 }
 
 impl CollabPlugin for IndexeddbDiskPlugin {
-  fn init(&self, object_id: &str, _origin: &CollabOrigin, doc: &Doc) {
+  fn init(&self, object_id: &str, _origin: &CollabOrigin, ctx: &mut CollabContext) {
     if let Some(db) = self.collab_db.upgrade() {
       let object_id = object_id.to_string();
-      let doc = doc.clone();
+      let doc = ctx.doc().clone();
       let uid = self.uid;
 
       tokio::task::spawn_local(async move {
         match db.load_doc(uid, &object_id, doc.clone()).await {
-          Ok(_) => {},
+          Ok(version) => {
+            ctx.version = version;
+          },
           Err(err) => {
             if err.is_record_not_found() {
               let encoded_collab = doc.get_encoded_collab_v1();

collab-plugins/src/local_storage/kv/db.rs

@@ -7,6 +7,7 @@ use crate::local_storage::kv::PersistenceError;
 use crate::local_storage::kv::keys::*;
 use crate::local_storage::kv::oid::{DocIDGen, OID};
 use crate::local_storage::kv::snapshot::CollabSnapshot;
+use collab::core::collab::CollabVersion;
 use smallvec::SmallVec;
 use yrs::{TransactionMut, Update};
 
@@ -112,6 +113,7 @@ pub fn insert_doc_update<'a, K, S>(
   db: &S,
   doc_id: DocID,
   object_id: &K,
+  version: Option<&CollabVersion>,
   value: Vec<u8>,
 ) -> Result<Vec<u8>, PersistenceError>
 where

collab-plugins/src/local_storage/kv/doc.rs

@@ -1,6 +1,7 @@
 use crate::local_storage::kv::keys::*;
 use crate::local_storage::kv::snapshot::SnapshotAction;
 use crate::local_storage::kv::*;
+use collab::core::collab::{CollabVersion, VersionedData};
 use smallvec::{SmallVec, smallvec};
 use std::collections::HashSet;
 use tracing::{error, info};
@@ -19,6 +20,7 @@ where
     uid: i64,
     workspace_id: &str,
     object_id: &str,
+    collab_version: Option<&CollabVersion>,
     txn: &T,
   ) -> Result<(), PersistenceError> {
     if self.is_exist(uid, workspace_id, object_id) {
@@ -28,7 +30,7 @@ where
     let doc_id = get_or_create_did(uid, self, workspace_id, object_id)?;
     let doc_state = txn.encode_diff_v1(&StateVector::default());
     let sv = txn.state_vector().encode_v1();
-    let doc_state_key = make_doc_state_key(doc_id);
+    let doc_state_key = make_doc_state_key(doc_id, collab_version);
     let sv_key = make_state_vector_key(doc_id);
 
     info!("new doc:{:?}, doc state len:{}", object_id, doc_state.len());
@@ -43,6 +45,7 @@ where
     uid: i64,
     workspace_id: &str,
     object_id: &str,
+    collab_version: Option<&CollabVersion>,
     state_vector: Vec<u8>,
     doc_state: Vec<u8>,
   ) -> Result<(), PersistenceError> {
@@ -51,7 +54,7 @@ where
     let end = make_doc_end_key(doc_id);
     self.remove_range(start.as_ref(), end.as_ref())?;
 
-    let doc_state_key = make_doc_state_key(doc_id);
+    let doc_state_key = make_doc_state_key(doc_id, collab_version);
     let sv_key = make_state_vector_key(doc_id);
 
     info!("new doc:{:?}, doc state len:{}", object_id, doc_state.len());
@@ -73,6 +76,7 @@ where
     uid: i64,
     workspace_id: &str,
     object_id: &str,
+    collab_version: Option<&CollabVersion>,
     state_vector: Vec<u8>,
     doc_state: Vec<u8>,
   ) -> Result<(), PersistenceError> {
@@ -83,7 +87,7 @@ where
     let end = make_doc_end_key(doc_id);
     self.remove_range(start.as_ref(), end.as_ref())?;
 
-    let doc_state_key = make_doc_state_key(doc_id);
+    let doc_state_key = make_doc_state_key(doc_id, collab_version);
     let sv_key = make_state_vector_key(doc_id);
     // Insert new doc state and state vector
     self.insert(doc_state_key, doc_state)?;
@@ -95,6 +99,35 @@ where
     get_doc_id(uid, self, workspace_id, object_id).is_some()
   }
 
+  fn get_doc_state(&self, doc_id: DocID) -> Result<Option<VersionedData>, PersistenceError> {
+    let doc_state_start = make_doc_start_key(doc_id);
+    let doc_state_end = make_doc_end_key(doc_id);
+    let mut cursor = self.range(doc_state_start.clone()..doc_state_end)?;
+    if let Some(entry) = cursor.next() {
+      let key = entry.key();
+      let value = entry.value();
+      if key.starts_with(doc_state_start.as_ref()) {
+        let version = if key.len() > doc_state_start.len() {
+          let collab_version: [u8; 16] =
+            match key[doc_state_start.len()..doc_state_start.len() + 16].try_into() {
+              Ok(v) => v,
+              Err(_) => {
+                return Err(PersistenceError::InvalidData(format!(
+                  "invalid collab version in doc state key: {:?}",
+                  key
+                )));
+              },
+            };
+          Some(CollabVersion::from_bytes(collab_version))
+        } else {
+          None
+        };
+        return Ok(Some(VersionedData::new(value, version)));
+      }
+    }
+    Ok(None)
+  }
+
   /// Load the document from the database and apply the updates to the transaction.
   /// It will try to load the document in these two ways:
   ///   1. D = document state + updates
@@ -107,15 +140,14 @@ where
     workspace_id: &str,
     object_id: &str,
     txn: &mut TransactionMut,
-  ) -> Result<u32, PersistenceError> {
+  ) -> Result<Option<CollabVersion>, PersistenceError> {
     let mut update_count = 0;
+    let mut collab_version = None;
 
     if let Some(doc_id) = get_doc_id(uid, self, workspace_id, object_id) {
-      let doc_state_key = make_doc_state_key(doc_id);
-      if let Some(doc_state) = self.get(doc_state_key.as_ref())? {
+      if let Some(versioned) = self.get_doc_state(doc_id)? {
         // Load the doc state
-
-        match Update::decode_v1(doc_state.as_ref()) {
+        match Update::decode_v1(&versioned.data) {
           Ok(update) => {
             txn.try_apply_update(update)?;
           },
@@ -125,6 +157,8 @@ where
           },
         }
 
+        collab_version = versioned.version;
+
         // If the enable_snapshot is true, we will try to load the snapshot.
         let update_start = make_doc_update_key(doc_id, 0).to_vec();
         let update_end = make_doc_update_key(doc_id, Clock::MAX);
@@ -155,7 +189,7 @@ where
         );
       }
 
-      Ok(update_count)
+      Ok(collab_version)
     } else {
       tracing::trace!("[Client] => {:?} not exist", object_id);
       Err(PersistenceError::RecordNotFound(format!(
@@ -171,7 +205,7 @@ where
     workspace_id: &str,
     object_id: &str,
     doc: &Doc,
-  ) -> Result<u32, PersistenceError> {
+  ) -> Result<Option<CollabVersion>, PersistenceError> {
     let mut txn = doc.transact_mut();
     self.load_doc_with_txn(uid, workspace_id, object_id, &mut txn)
   }
@@ -182,6 +216,7 @@ where
     uid: i64,
     workspace_id: &str,
     object_id: &str,
+    version: Option<&CollabVersion>,
     update: &[u8],
   ) -> Result<Vec<u8>, PersistenceError> {
     match get_doc_id(uid, self, workspace_id, object_id) {
@@ -196,7 +231,7 @@ where
           object_id
         )))
       },
-      Some(doc_id) => insert_doc_update(self, doc_id, object_id, update.to_vec()),
+      Some(doc_id) => insert_doc_update(self, doc_id, object_id, version, update.to_vec()),
     }
   }
 
@@ -234,6 +269,7 @@ where
     uid: i64,
     workspace_id: &str,
     object_id: &str,
+    version: Option<&CollabVersion>,
     doc_state: &[u8],
     sv: &[u8],
   ) -> Result<(), PersistenceError> {
@@ -242,7 +278,7 @@ where
     let end = make_doc_end_key(doc_id);
     self.remove_range(start.as_ref(), end.as_ref())?;
 
-    let doc_state_key = make_doc_state_key(doc_id);
+    let doc_state_key = make_doc_state_key(doc_id, version);
     let sv_key = make_state_vector_key(doc_id);
 
     // Insert new doc state and state vector
@@ -294,10 +330,9 @@ where
       self.remove_range(start.as_ref(), end.as_ref())?;
 
       // Delete the document state and the state vector
-      let doc_state_key = make_doc_state_key(did);
-      let sv_key = make_state_vector_key(did);
-      let _ = self.remove(doc_state_key.as_ref());
-      let _ = self.remove(sv_key.as_ref());
+      let doc_state_key = make_doc_start_key(did);
+      let sv_key = make_doc_end_key(did);
+      let _ = self.remove_range(doc_state_key.as_ref(), sv_key.as_ref());
 
       // Delete the snapshot
       self.delete_all_snapshots(uid, object_id)?;

collab-plugins/src/local_storage/kv/keys.rs

@@ -1,8 +1,8 @@
 use std::io::Write;
 use std::ops::Deref;
 
+use collab::core::collab::CollabVersion;
 use smallvec::{SmallVec, smallvec};
-
 // https://github.com/spacejam/sled
 // sled performs prefix encoding on long keys with similar prefixes that are grouped together in a
 // range, as well as suffix truncation to further reduce the indexing costs of long keys. Nodes
@@ -99,17 +99,24 @@ pub fn oid_from_key(key: &[u8]) -> &[u8] {
   &key[10..(key.len() - 1)]
 }
 
-// [1,1,  0,0,0,0,0,0,0,0,  0]
-pub fn make_doc_state_key(doc_id: DocID) -> Key<DOC_STATE_KEY_LEN> {
+// [1,1,  0,0,0,0,0,0,0,0,  0] if version is None
+// [1,1,  0,0,0,0,0,0,0,0,  0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] if version is Some
+pub fn make_doc_state_key(
+  doc_id: DocID,
+  version: Option<&CollabVersion>,
+) -> Key<DOC_STATE_KEY_LEN> {
   let mut v: SmallVec<[u8; DOC_STATE_KEY_LEN]> = smallvec![DOC_SPACE, DOC_SPACE_OBJECT_KEY];
   v.write_all(&doc_id.to_be_bytes()).unwrap();
   v.push(DOC_STATE);
+  if let Some(ver) = version {
+    v.write_all(ver.as_bytes()).unwrap();
+  }
   Key(v)
 }
 
 // document related elements are stored within bounds [0,1,..did,0]..[0,1,..did,255]
 pub fn make_doc_start_key(doc_id: DocID) -> Key<DOC_STATE_KEY_LEN> {
-  make_doc_state_key(doc_id)
+  make_doc_state_key(doc_id, None)
 }
 // [1,1,  0,0,0,0,0,0,0,0,  255]
 pub fn make_doc_end_key(doc_id: DocID) -> Key<DOC_STATE_KEY_LEN> {

collab-plugins/src/local_storage/rocksdb/rocksdb_plugin.rs

@@ -94,9 +94,16 @@ impl RocksdbDiskPlugin {
       if !rocksdb_read.is_exist(self.uid, &self.workspace_id, &self.object_id) {
         match self.collab_type.validate_require_data(collab) {
           Ok(_) => {
+            let version = collab.version();
             let txn = collab.transact();
             if let Err(err) = collab_db.with_write_txn(|w_db_txn| {
-              w_db_txn.create_new_doc(self.uid, &self.workspace_id, &self.object_id, &txn)?;
+              w_db_txn.create_new_doc(
+                self.uid,
+                &self.workspace_id,
+                &self.object_id,
+                version,
+                &txn,
+              )?;
               info!(
                 "[Rocksdb Plugin]: created new doc {}, collab_type:{}",
                 self.object_id, self.collab_type
@@ -132,7 +139,7 @@ impl CollabPlugin for RocksdbDiskPlugin {
     object_id: &str,
     _txn: &TransactionMut,
     update: &[u8],
-    _collab_version: Option<&CollabVersion>,
+    collab_version: Option<&CollabVersion>,
   ) {
     // Only push update if the doc is loaded
     if !self.did_init.load(SeqCst) {
@@ -142,7 +149,13 @@ impl CollabPlugin for RocksdbDiskPlugin {
       self.increase_count();
       //Acquire a write transaction to ensure consistency
       let result = db.with_write_txn(|w_db_txn| {
-        let _ = w_db_txn.push_update(self.uid, self.workspace_id.as_str(), object_id, update)?;
+        let _ = w_db_txn.push_update(
+          self.uid,
+          self.workspace_id.as_str(),
+          object_id,
+          collab_version,
+          update,
+        )?;
         use yrs::updates::decoder::Decode;
         tracing::trace!(
           "[Rocksdb Plugin]: Collab {} {} persisting update: {:#?}",

collab-plugins/tests/disk/insert_test.rs

@@ -75,6 +75,7 @@ async fn flush_test() {
       test.uid,
       &test.workspace_id,
       &object_id_str,
+      encode_collab.collab_version.as_ref(),
       encode_collab.state_vector.to_vec(),
       encode_collab.doc_state.to_vec(),
     )