AppFlowy/frontend/rust-lib/flowy-document/src/manager.rs

use crate::editor::{initial_document_content, AppFlowyDocumentEditor, DocumentRevisionMergeable};
use crate::entities::{DocumentVersionPB, EditParams};
use crate::old_editor::editor::{DeltaDocumentEditor, DeltaDocumentRevisionMergeable};
use crate::old_editor::snapshot::DeltaDocumentSnapshotPersistence;
use crate::services::rev_sqlite::{
  SQLiteDeltaDocumentRevisionPersistence, SQLiteDocumentRevisionPersistence,
  SQLiteDocumentRevisionSnapshotPersistence,
};
use crate::services::DocumentPersistence;
use crate::{errors::FlowyError, DocumentCloudService};
use bytes::Bytes;
use document_model::document::DocumentId;
use flowy_client_sync::client_document::initial_delta_document_content;
use flowy_error::FlowyResult;
use flowy_revision::{
  RevisionCloudService, RevisionManager, RevisionPersistence, RevisionPersistenceConfiguration,
  RevisionWebSocket,
};
use flowy_sqlite::ConnectionPool;
use lib_infra::async_trait::async_trait;
use lib_infra::future::FutureResult;
use lib_infra::ref_map::{RefCountHashMap, RefCountValue};
use lib_infra::util::md5;
use lib_ws::WSConnectState;
use revision_model::Revision;
use std::any::Any;
use std::convert::TryFrom;
use std::sync::Arc;
use tokio::sync::RwLock;
use ws_model::ws_revision::ServerRevisionWSData;

pub trait DocumentUser: Send + Sync {
  fn user_dir(&self) -> Result<String, FlowyError>;
  fn user_id(&self) -> Result<String, FlowyError>;
  fn token(&self) -> Result<String, FlowyError>;
}

pub trait DocumentDatabase: Send + Sync {
  fn db_pool(&self) -> Result<Arc<ConnectionPool>, FlowyError>;
}

#[async_trait]
pub trait DocumentEditor: Send + Sync {
  /// Called when the document get closed
  async fn close(&self);

  /// Exports the document content. The content is encoded in the corresponding
  /// editor data format.
  fn export(&self) -> FutureResult<String, FlowyError>;

  /// Duplicate the document inner data into String
  fn duplicate(&self) -> FutureResult<String, FlowyError>;

  fn receive_ws_data(&self, data: ServerRevisionWSData) -> FutureResult<(), FlowyError>;

  fn receive_ws_state(&self, state: &WSConnectState);

  /// Receives the local operations made by the user input. The operations are encoded
  /// in binary format.
  fn compose_local_operations(&self, data: Bytes) -> FutureResult<(), FlowyError>;

  /// Returns the `Any` reference that can be used to downcast back to the original,
  /// concrete type.
  ///
  /// The indirection through `as_any` is because using `downcast_ref`
  /// on `Box<A>` *directly* only lets us downcast back to `&A` again. You can take a look at [this](https://stackoverflow.com/questions/33687447/how-to-get-a-reference-to-a-concrete-type-from-a-trait-object)
  /// for more information.
  ///
  ///
  fn as_any(&self) -> &dyn Any;
}

#[derive(Clone, Debug)]
pub struct DocumentConfig {
  pub version: DocumentVersionPB,
}

impl std::default::Default for DocumentConfig {
  fn default() -> Self {
    Self {
      version: DocumentVersionPB::V1,
    }
  }
}

pub struct DocumentManager {
  cloud_service: Arc<dyn DocumentCloudService>,
  rev_web_socket: Arc<dyn RevisionWebSocket>,
  editor_map: Arc<RwLock<RefCountHashMap<RefCountDocumentHandler>>>,
  user: Arc<dyn DocumentUser>,
  persistence: Arc<DocumentPersistence>,
  #[allow(dead_code)]
  config: DocumentConfig,
}

impl DocumentManager {
  pub fn new(
    cloud_service: Arc<dyn DocumentCloudService>,
    document_user: Arc<dyn DocumentUser>,
    database: Arc<dyn DocumentDatabase>,
    rev_web_socket: Arc<dyn RevisionWebSocket>,
    config: DocumentConfig,
  ) -> Self {
    Self {
      cloud_service,
      rev_web_socket,
      editor_map: Arc::new(RwLock::new(RefCountHashMap::new())),
      user: document_user,
      persistence: Arc::new(DocumentPersistence::new(database)),
      config,
    }
  }

  /// Called immediately after the application launched with the user sign in/sign up.
  #[tracing::instrument(level = "trace", skip_all, err)]
  pub async fn initialize(&self, user_id: &str) -> FlowyResult<()> {
    self.persistence.initialize(user_id)?;
    listen_ws_state_changed(self.rev_web_socket.clone(), self.editor_map.clone());
    Ok(())
  }

  pub async fn initialize_with_new_user(&self, _user_id: &str, _token: &str) -> FlowyResult<()> {
    Ok(())
  }

  #[tracing::instrument(level = "trace", skip_all, fields(document_id), err)]
  pub async fn open_document_editor<T: AsRef<str>>(
    &self,
    document_id: T,
  ) -> Result<Arc<dyn DocumentEditor>, FlowyError> {
    let document_id = document_id.as_ref();
    tracing::Span::current().record("document_id", document_id);
    self.init_document_editor(document_id).await
  }

  #[tracing::instrument(level = "trace", skip(self, editor_id), fields(editor_id), err)]
  pub async fn close_document_editor<T: AsRef<str>>(&self, editor_id: T) -> Result<(), FlowyError> {
    let editor_id = editor_id.as_ref();
    tracing::Span::current().record("editor_id", editor_id);
    self.editor_map.write().await.remove(editor_id).await;
    Ok(())
  }

  pub async fn apply_edit(&self, params: EditParams) -> FlowyResult<()> {
    let editor = self.get_document_editor(&params.doc_id).await?;
    editor
      .compose_local_operations(Bytes::from(params.operations))
      .await?;
    Ok(())
  }

  pub async fn create_document<T: AsRef<str>>(
    &self,
    doc_id: T,
    revisions: Vec<Revision>,
  ) -> FlowyResult<()> {
    let doc_id = doc_id.as_ref().to_owned();
    let db_pool = self.persistence.database.db_pool()?;
    // Maybe we could save the document to disk without creating the RevisionManager
    let rev_manager = self.make_rev_manager(&doc_id, db_pool)?;
    rev_manager.reset_object(revisions).await?;
    Ok(())
  }

  pub async fn receive_ws_data(&self, data: Bytes) {
    let result: Result<ServerRevisionWSData, serde_json::Error> =
      ServerRevisionWSData::try_from(data);
    match result {
      Ok(data) => match self.editor_map.read().await.get(&data.object_id) {
        None => tracing::error!(
          "Can't find any source handler for {:?}-{:?}",
          data.object_id,
          data.payload
        ),
        Some(handler) => match handler.0.receive_ws_data(data).await {
          Ok(_) => {},
          Err(e) => tracing::error!("{}", e),
        },
      },
      Err(e) => {
        tracing::error!("Document ws data parser failed: {:?}", e);
      },
    }
  }

  pub fn initial_document_content(&self) -> String {
    match self.config.version {
      DocumentVersionPB::V0 => initial_delta_document_content(),
      DocumentVersionPB::V1 => initial_document_content(),
    }
  }
}

impl DocumentManager {
  /// Returns the `DocumentEditor`
  ///
  /// # Arguments
  ///
  /// * `doc_id`: the id of the document
  ///
  /// returns: Result<Arc<DocumentEditor>, FlowyError>
  ///
  async fn get_document_editor(&self, doc_id: &str) -> FlowyResult<Arc<dyn DocumentEditor>> {
    match self.editor_map.read().await.get(doc_id) {
      None => {
        //
        tracing::warn!("Should call init_document_editor first");
        self.init_document_editor(doc_id).await
      },
      Some(handler) => Ok(handler.0.clone()),
    }
  }

  /// Initializes a document editor with the doc_id
  ///
  /// # Arguments
  ///
  /// * `doc_id`: the id of the document
  /// * `pool`: sqlite connection pool
  ///
  /// returns: Result<Arc<DocumentEditor>, FlowyError>
  ///
  #[tracing::instrument(level = "trace", skip(self), err)]
  pub async fn init_document_editor(
    &self,
    doc_id: &str,
  ) -> Result<Arc<dyn DocumentEditor>, FlowyError> {
    let pool = self.persistence.database.db_pool()?;
    let user = self.user.clone();
    let token = self.user.token()?;
    let cloud_service = Arc::new(DocumentRevisionCloudService {
      token,
      server: self.cloud_service.clone(),
    });

    match self.config.version {
      DocumentVersionPB::V0 => {
        let rev_manager = self.make_delta_document_rev_manager(doc_id, pool.clone())?;
        let editor: Arc<dyn DocumentEditor> = Arc::new(
          DeltaDocumentEditor::new(
            doc_id,
            user,
            rev_manager,
            self.rev_web_socket.clone(),
            cloud_service,
          )
          .await?,
        );
        self
          .editor_map
          .write()
          .await
          .insert(doc_id.to_string(), RefCountDocumentHandler(editor.clone()));
        Ok(editor)
      },
      DocumentVersionPB::V1 => {
        let rev_manager = self.make_document_rev_manager(doc_id, pool.clone())?;
        let editor: Arc<dyn DocumentEditor> =
          Arc::new(AppFlowyDocumentEditor::new(doc_id, user, rev_manager, cloud_service).await?);
        self
          .editor_map
          .write()
          .await
          .insert(doc_id.to_string(), RefCountDocumentHandler(editor.clone()));
        Ok(editor)
      },
    }
  }

  fn make_rev_manager(
    &self,
    doc_id: &str,
    pool: Arc<ConnectionPool>,
  ) -> Result<RevisionManager<Arc<ConnectionPool>>, FlowyError> {
    match self.config.version {
      DocumentVersionPB::V0 => self.make_delta_document_rev_manager(doc_id, pool),
      DocumentVersionPB::V1 => self.make_document_rev_manager(doc_id, pool),
    }
  }

  fn make_document_rev_manager(
    &self,
    doc_id: &str,
    pool: Arc<ConnectionPool>,
  ) -> Result<RevisionManager<Arc<ConnectionPool>>, FlowyError> {
    let user_id = self.user.user_id()?;
    let disk_cache = SQLiteDocumentRevisionPersistence::new(&user_id, pool.clone());
    let configuration = RevisionPersistenceConfiguration::new(200, true);
    let rev_persistence = RevisionPersistence::new(&user_id, doc_id, disk_cache, configuration);
    let snapshot_persistence = SQLiteDocumentRevisionSnapshotPersistence::new(doc_id, pool);
    Ok(RevisionManager::new(
      &user_id,
      doc_id,
      rev_persistence,
      DocumentRevisionMergeable(),
      snapshot_persistence,
    ))
  }

  fn make_delta_document_rev_manager(
    &self,
    doc_id: &str,
    pool: Arc<ConnectionPool>,
  ) -> Result<RevisionManager<Arc<ConnectionPool>>, FlowyError> {
    let user_id = self.user.user_id()?;
    let disk_cache = SQLiteDeltaDocumentRevisionPersistence::new(&user_id, pool);
    let configuration = RevisionPersistenceConfiguration::new(100, true);
    let rev_persistence = RevisionPersistence::new(&user_id, doc_id, disk_cache, configuration);
    Ok(RevisionManager::new(
      &user_id,
      doc_id,
      rev_persistence,
      DeltaDocumentRevisionMergeable(),
      DeltaDocumentSnapshotPersistence(),
    ))
  }
}

struct DocumentRevisionCloudService {
  token: String,
  server: Arc<dyn DocumentCloudService>,
}

impl RevisionCloudService for DocumentRevisionCloudService {
  #[tracing::instrument(level = "trace", skip(self))]
  fn fetch_object(
    &self,
    user_id: &str,
    object_id: &str,
  ) -> FutureResult<Vec<Revision>, FlowyError> {
    let params: DocumentId = object_id.to_string().into();
    let server = self.server.clone();
    let token = self.token.clone();

    FutureResult::new(async move {
      match server.fetch_document(&token, params).await? {
        None => Err(FlowyError::record_not_found().context("Remote doesn't have this document")),
        Some(payload) => {
          let bytes = Bytes::from(payload.data.clone());
          let doc_md5 = md5(&bytes);
          let revision = Revision::new(
            &payload.doc_id,
            payload.base_rev_id,
            payload.rev_id,
            bytes,
            doc_md5,
          );
          Ok(vec![revision])
        },
      }
    })
  }
}

#[derive(Clone)]
struct RefCountDocumentHandler(Arc<dyn DocumentEditor>);

#[async_trait]
impl RefCountValue for RefCountDocumentHandler {
  async fn did_remove(&self) {
    self.0.close().await;
  }
}

impl std::ops::Deref for RefCountDocumentHandler {
  type Target = Arc<dyn DocumentEditor>;

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

#[tracing::instrument(level = "trace", skip(web_socket, handlers))]
fn listen_ws_state_changed(
  web_socket: Arc<dyn RevisionWebSocket>,
  handlers: Arc<RwLock<RefCountHashMap<RefCountDocumentHandler>>>,
) {
  tokio::spawn(async move {
    let mut notify = web_socket.subscribe_state_changed().await;
    while let Ok(state) = notify.recv().await {
      handlers.read().await.values().iter().for_each(|handler| {
        handler.receive_ws_state(&state);
      })
    }
  });
}