AppFlowy/frontend/rust-lib/flowy-revision/src/rev_snapshot.rs

#![allow(clippy::all)]
#![allow(dead_code)]
#![allow(unused_variables)]
use crate::{RevIdCounter, RevisionMergeable, RevisionObjectDeserializer, RevisionPersistence};
use bytes::Bytes;
use flowy_error::FlowyResult;
use revision_model::Revision;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::Arc;

pub trait RevisionSnapshotPersistence: Send + Sync {
  fn should_generate_snapshot_from_range(&self, start_rev_id: i64, current_rev_id: i64) -> bool {
    (current_rev_id - start_rev_id) >= AUTO_GEN_SNAPSHOT_PER_10_REVISION
  }

  fn write_snapshot(&self, rev_id: i64, data: Vec<u8>) -> FlowyResult<()>;

  fn read_snapshot(&self, rev_id: i64) -> FlowyResult<Option<RevisionSnapshotData>>;

  fn read_last_snapshot(&self) -> FlowyResult<Option<RevisionSnapshotData>>;
}

pub trait RevisionSnapshotDataGenerator: Send + Sync {
  fn generate_snapshot_data(&self) -> Option<RevisionSnapshotData>;
}

const AUTO_GEN_SNAPSHOT_PER_10_REVISION: i64 = 10;

pub struct RevisionSnapshotController<Connection> {
  user_id: String,
  object_id: String,
  rev_snapshot_persistence: Arc<dyn RevisionSnapshotPersistence>,
  rev_snapshot_data: Option<Arc<dyn RevisionSnapshotDataGenerator>>,
  rev_id_counter: Arc<RevIdCounter>,
  rev_persistence: Arc<RevisionPersistence<Connection>>,
  rev_compress: Arc<dyn RevisionMergeable>,
  start_rev_id: AtomicI64,
}

impl<Connection> RevisionSnapshotController<Connection>
where
  Connection: 'static,
{
  pub fn new<D>(
    user_id: &str,
    object_id: &str,
    disk_cache: D,
    rev_id_counter: Arc<RevIdCounter>,
    revision_persistence: Arc<RevisionPersistence<Connection>>,
    revision_compress: Arc<dyn RevisionMergeable>,
  ) -> Self
  where
    D: RevisionSnapshotPersistence + 'static,
  {
    let rev_snapshot_persistence = Arc::new(disk_cache);
    Self {
      user_id: user_id.to_string(),
      object_id: object_id.to_string(),
      rev_snapshot_persistence,
      rev_id_counter,
      start_rev_id: AtomicI64::new(0),
      rev_snapshot_data: None,
      rev_persistence: revision_persistence,
      rev_compress: revision_compress,
    }
  }

  pub async fn set_snapshot_data_generator(
    &mut self,
    generator: Arc<dyn RevisionSnapshotDataGenerator>,
  ) {
    self.rev_snapshot_data = Some(generator);
  }

  pub async fn generate_snapshot(&self) {
    if let Some((rev_id, bytes)) = self.generate_snapshot_data() {
      if let Err(e) = self
        .rev_snapshot_persistence
        .write_snapshot(rev_id, bytes.to_vec())
      {
        tracing::error!("Save snapshot failed: {}", e);
      }
    }
  }

  /// Find the nearest revision base on the passed-in rev_id
  #[tracing::instrument(level = "trace", skip_all)]
  pub fn restore_from_snapshot<B>(&self, rev_id: i64) -> Option<(B::Output, Revision)>
  where
    B: RevisionObjectDeserializer,
  {
    tracing::info!("[Restore] Try to find if {} has snapshot", self.object_id);
    let snapshot = self.rev_snapshot_persistence.read_last_snapshot().ok()??;
    let snapshot_rev_id = snapshot.rev_id;
    let revision = Revision::new(
      &self.object_id,
      snapshot.base_rev_id,
      snapshot.rev_id,
      snapshot.data,
      "".to_owned(),
    );
    tracing::info!(
      "[Restore] Try to restore from snapshot: {}, {}",
      snapshot.base_rev_id,
      snapshot.rev_id
    );
    let object = B::deserialize_revisions(&self.object_id, vec![revision.clone()]).ok()?;
    tracing::info!(
      "[Restore] Restore {} from snapshot with rev_id: {}",
      self.object_id,
      snapshot_rev_id
    );

    Some((object, revision))
  }

  pub fn generate_snapshot_if_need(&self) {
    let current_rev_id = self.rev_id_counter.value();
    let start_rev_id = self.get_start_rev_id();
    if current_rev_id <= start_rev_id {
      return;
    }
    if self
      .rev_snapshot_persistence
      .should_generate_snapshot_from_range(start_rev_id, current_rev_id)
    {
      if let Some((rev_id, bytes)) = self.generate_snapshot_data() {
        let disk_cache = self.rev_snapshot_persistence.clone();
        tokio::spawn(async move {
          let _ = disk_cache.write_snapshot(rev_id, bytes.to_vec());
        });
      }
      self.set_start_rev_id(current_rev_id);
    }
  }

  fn generate_snapshot_data(&self) -> Option<(i64, Bytes)> {
    let revisions = self
      .rev_persistence
      .load_all_records(&self.object_id)
      .map(|records| {
        records
          .into_iter()
          .map(|record| record.revision)
          .collect::<Vec<Revision>>()
      })
      .ok()?;

    if revisions.is_empty() {
      return None;
    }

    let data = self.rev_compress.combine_revisions(revisions).ok()?;
    let rev_id = self.rev_id_counter.value();
    Some((rev_id, data))
  }

  fn get_start_rev_id(&self) -> i64 {
    self.start_rev_id.load(SeqCst)
  }

  fn set_start_rev_id(&self, rev_id: i64) {
    let _ = self
      .start_rev_id
      .fetch_update(SeqCst, SeqCst, |_| Some(rev_id));
  }
}

impl<Connection> std::ops::Deref for RevisionSnapshotController<Connection> {
  type Target = Arc<dyn RevisionSnapshotPersistence>;

  fn deref(&self) -> &Self::Target {
    &self.rev_snapshot_persistence
  }
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct RevisionSnapshotData {
  pub rev_id: i64,
  pub base_rev_id: i64,
  pub timestamp: i64,
  pub data: Bytes,
}