engine.rs

//! Handlers and actors to for live syncing replicas.
//!
//! [`crate::Replica`] is also called documents here.

use std::{io, path::PathBuf, str::FromStr, sync::Arc};

use anyhow::{bail, ensure, Context, Result};
use futures_lite::{Stream, StreamExt};
use iroh::{key::PublicKey, Endpoint, NodeAddr};
use iroh_blobs::{
    downloader::Downloader, net_protocol::ProtectCb, store::EntryStatus,
    util::local_pool::LocalPoolHandle, Hash,
};
use iroh_gossip::net::Gossip;
use serde::{Deserialize, Serialize};
use tokio::sync::{mpsc, oneshot, RwLock};
use tokio_util::task::AbortOnDropHandle;
use tracing::{error, error_span, Instrument};

use self::live::{LiveActor, ToLiveActor};
pub use self::{
    live::SyncEvent,
    state::{Origin, SyncReason},
};
use crate::{
    actor::SyncHandle, Author, AuthorId, ContentStatus, ContentStatusCallback, Entry, NamespaceId,
};

mod gossip;
mod live;
mod state;

/// Capacity of the channel for the [`ToLiveActor`] messages.
const ACTOR_CHANNEL_CAP: usize = 64;
/// Capacity for the channels for [`Engine::subscribe`].
const SUBSCRIBE_CHANNEL_CAP: usize = 256;

/// The sync engine coordinates actors that manage open documents, set-reconciliation syncs with
/// peers and a gossip swarm for each syncing document.
#[derive(derive_more::Debug, Clone)]
pub struct Engine<D> {
    /// [`Endpoint`] used by the engine.
    pub(crate) endpoint: Endpoint,
    /// Handle to the actor thread.
    pub(crate) sync: SyncHandle,
    /// The persistent default author for this engine.
    default_author: Arc<DefaultAuthor>,
    to_live_actor: mpsc::Sender<ToLiveActor>,
    #[allow(dead_code)]
    actor_handle: Arc<AbortOnDropHandle<()>>,
    #[debug("ContentStatusCallback")]
    content_status_cb: ContentStatusCallback,
    local_pool_handle: LocalPoolHandle,
    blob_store: D,
    #[cfg(feature = "rpc")]
    pub(crate) rpc_handler: Arc<std::sync::OnceLock<crate::rpc::RpcHandler>>,
}

impl<D: iroh_blobs::store::Store> Engine<D> {
    /// Start the sync engine.
    ///
    /// This will spawn two tokio tasks for the live sync coordination and gossip actors, and a
    /// thread for the [`crate::actor::SyncHandle`].
    pub async fn spawn(
        endpoint: Endpoint,
        gossip: Gossip,
        replica_store: crate::store::Store,
        bao_store: D,
        downloader: Downloader,
        default_author_storage: DefaultAuthorStorage,
        local_pool_handle: LocalPoolHandle,
    ) -> anyhow::Result<Self> {
        let (live_actor_tx, to_live_actor_recv) = mpsc::channel(ACTOR_CHANNEL_CAP);
        let me = endpoint.node_id().fmt_short();

        let content_status_cb = {
            let bao_store = bao_store.clone();
            Arc::new(move |hash| entry_to_content_status(bao_store.entry_status_sync(&hash)))
        };
        let sync = SyncHandle::spawn(replica_store, Some(content_status_cb.clone()), me.clone());

        let actor = LiveActor::new(
            sync.clone(),
            endpoint.clone(),
            gossip.clone(),
            bao_store.clone(),
            downloader,
            to_live_actor_recv,
            live_actor_tx.clone(),
        );
        let actor_handle = tokio::task::spawn(
            async move {
                if let Err(err) = actor.run().await {
                    error!("sync actor failed: {err:?}");
                }
            }
            .instrument(error_span!("sync", %me)),
        );

        let default_author = match DefaultAuthor::load(default_author_storage, &sync).await {
            Ok(author) => author,
            Err(err) => {
                // If loading the default author failed, make sure to shutdown the sync actor before
                // returning.
                let _store = sync.shutdown().await.ok();
                return Err(err);
            }
        };

        Ok(Self {
            endpoint,
            sync,
            to_live_actor: live_actor_tx,
            actor_handle: Arc::new(AbortOnDropHandle::new(actor_handle)),
            content_status_cb,
            default_author: Arc::new(default_author),
            local_pool_handle,
            blob_store: bao_store,
            #[cfg(feature = "rpc")]
            rpc_handler: Default::default(),
        })
    }

    /// Return a callback that can be added to blobs to protect the content of
    /// all docs from garbage collection.
    pub fn protect_cb(&self) -> ProtectCb {
        let this = self.clone();
        Box::new(move |live| {
            let this = this.clone();
            Box::pin(async move {
                let doc_hashes = match this.sync.content_hashes().await {
                    Ok(hashes) => hashes,
                    Err(err) => {
                        tracing::warn!("Error getting doc hashes: {}", err);
                        return;
                    }
                };
                for hash in doc_hashes {
                    match hash {
                        Ok(hash) => {
                            live.insert(hash);
                        }
                        Err(err) => {
                            tracing::error!("Error getting doc hash: {}", err);
                        }
                    }
                }
            })
        })
    }

    /// Get the blob store.
    pub fn blob_store(&self) -> &D {
        &self.blob_store
    }

    /// Start to sync a document.
    ///
    /// If `peers` is non-empty, it will both do an initial set-reconciliation sync with each peer,
    /// and join an iroh-gossip swarm with these peers to receive and broadcast document updates.
    pub async fn start_sync(&self, namespace: NamespaceId, peers: Vec<NodeAddr>) -> Result<()> {
        let (reply, reply_rx) = oneshot::channel();
        self.to_live_actor
            .send(ToLiveActor::StartSync {
                namespace,
                peers,
                reply,
            })
            .await?;
        reply_rx.await??;
        Ok(())
    }

    /// Stop the live sync for a document and leave the gossip swarm.
    ///
    /// If `kill_subscribers` is true, all existing event subscribers will be dropped. This means
    /// they will receive `None` and no further events in case of rejoining the document.
    pub async fn leave(&self, namespace: NamespaceId, kill_subscribers: bool) -> Result<()> {
        let (reply, reply_rx) = oneshot::channel();
        self.to_live_actor
            .send(ToLiveActor::Leave {
                namespace,
                kill_subscribers,
                reply,
            })
            .await?;
        reply_rx.await??;
        Ok(())
    }

    /// Subscribe to replica and sync progress events.
    pub async fn subscribe(
        &self,
        namespace: NamespaceId,
    ) -> Result<impl Stream<Item = Result<LiveEvent>> + Unpin + 'static> {
        let content_status_cb = self.content_status_cb.clone();

        // Create a future that sends channel senders to the respective actors.
        // We clone `self` so that the future does not capture any lifetimes.
        let this = self.clone();

        // Subscribe to insert events from the replica.
        let a = {
            let (s, r) = async_channel::bounded(SUBSCRIBE_CHANNEL_CAP);
            this.sync.subscribe(namespace, s).await?;
            Box::pin(r).map(move |ev| LiveEvent::from_replica_event(ev, &content_status_cb))
        };

        // Subscribe to events from the [`live::Actor`].
        let b = {
            let (s, r) = async_channel::bounded(SUBSCRIBE_CHANNEL_CAP);
            let r = Box::pin(r);
            let (reply, reply_rx) = oneshot::channel();
            this.to_live_actor
                .send(ToLiveActor::Subscribe {
                    namespace,
                    sender: s,
                    reply,
                })
                .await?;
            reply_rx.await??;
            r.map(|event| Ok(LiveEvent::from(event)))
        };

        Ok(a.or(b))
    }

    /// Handle an incoming iroh-docs connection.
    pub async fn handle_connection(&self, conn: iroh::endpoint::Connecting) -> anyhow::Result<()> {
        self.to_live_actor
            .send(ToLiveActor::HandleConnection { conn })
            .await?;
        Ok(())
    }

    /// Shutdown the engine.
    pub async fn shutdown(&self) -> Result<()> {
        let (reply, reply_rx) = oneshot::channel();
        self.to_live_actor
            .send(ToLiveActor::Shutdown { reply })
            .await?;
        reply_rx.await?;
        Ok(())
    }

    /// Returns the stored `LocalPoolHandle`.
    pub fn local_pool_handle(&self) -> &LocalPoolHandle {
        &self.local_pool_handle
    }

    /// Authors API.
    pub fn authors(&self) -> Authors {
        Authors {
            sync: self.sync.clone(),
            default_author: self.default_author.clone(),
        }
    }
}

/// Authors client
#[derive(Debug, Clone)]
pub struct Authors {
    sync: SyncHandle,
    default_author: Arc<DefaultAuthor>,
}

impl Authors {
    /// Creates a new document author.
    ///
    /// You likely want to save the returned [`AuthorId`] somewhere so that you can use this author
    /// again.
    ///
    /// If you need only a single author, use [`Self::default`].
    pub async fn create(&self) -> Result<AuthorId> {
        let mut rng = rand::rngs::OsRng::default();
        let author = Author::new(&mut rng);
        self.sync.import_author(author).await
    }

    /// Returns the default document author of this node.
    ///
    /// On persistent nodes, the author is created on first start and its public key is saved
    /// in the data directory.
    ///
    /// The default author can be set with [`Self::set_default`].
    pub async fn default(&self) -> AuthorId {
        self.default_author.get().await
    }

    /// Sets the node-wide default author.
    ///
    /// If the author does not exist, an error is returned.
    ///
    /// On a persistent node, the author id will be saved to a file in the data directory and
    /// reloaded after a restart.
    pub async fn set_default(&self, author_id: AuthorId) -> Result<()> {
        self.default_author.set(author_id, &self.sync).await
    }

    /// Lists document authors for which we have a secret key.
    ///
    /// It's only possible to create writes from authors that we have the secret key of.
    pub async fn list(&self) -> Result<impl Stream<Item = Result<AuthorId>>> {
        let (tx, rx) = async_channel::bounded(64);
        self.sync.list_authors(tx).await?;

        Ok(rx)
    }

    /// Exports the given author.
    ///
    /// Warning: The [`Author`] struct contains sensitive data.
    pub async fn export(&self, author: AuthorId) -> Result<Option<Author>> {
        self.sync.export_author(author).await
    }

    /// Imports the given author.
    ///
    /// Warning: The [`Author`] struct contains sensitive data.
    pub async fn import(&self, author: Author) -> Result<AuthorId> {
        self.sync.import_author(author).await
    }

    /// Deletes the given author by id.
    ///
    /// Warning: This permanently removes this author.
    ///
    /// Returns an error if attempting to delete the default author.
    pub async fn delete(&self, author: AuthorId) -> Result<()> {
        let default_author = self.default().await;
        ensure!(
            author != default_author,
            "Deleting the default author is not supported"
        );
        self.sync.delete_author(author).await
    }
}

/// Converts an [`EntryStatus`] into a ['ContentStatus'].
pub fn entry_to_content_status(entry: io::Result<EntryStatus>) -> ContentStatus {
    match entry {
        Ok(EntryStatus::Complete) => ContentStatus::Complete,
        Ok(EntryStatus::Partial) => ContentStatus::Incomplete,
        Ok(EntryStatus::NotFound) => ContentStatus::Missing,
        Err(cause) => {
            tracing::warn!("Error while checking entry status: {cause:?}");
            ContentStatus::Missing
        }
    }
}

/// Events informing about actions of the live sync progress.
#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, strum::Display)]
pub enum LiveEvent {
    /// A local insertion.
    InsertLocal {
        /// The inserted entry.
        entry: Entry,
    },
    /// Received a remote insert.
    InsertRemote {
        /// The peer that sent us the entry.
        from: PublicKey,
        /// The inserted entry.
        entry: Entry,
        /// If the content is available at the local node
        content_status: ContentStatus,
    },
    /// The content of an entry was downloaded and is now available at the local node
    ContentReady {
        /// The content hash of the newly available entry content
        hash: Hash,
    },
    /// All pending content is now ready.
    ///
    /// This event signals that all queued content downloads from the last sync run have either
    /// completed or failed.
    ///
    /// It will only be emitted after a [`Self::SyncFinished`] event, never before.
    ///
    /// Receiving this event does not guarantee that all content in the document is available. If
    /// blobs failed to download, this event will still be emitted after all operations completed.
    PendingContentReady,
    /// We have a new neighbor in the swarm.
    NeighborUp(PublicKey),
    /// We lost a neighbor in the swarm.
    NeighborDown(PublicKey),
    /// A set-reconciliation sync finished.
    SyncFinished(SyncEvent),
}

impl From<live::Event> for LiveEvent {
    fn from(ev: live::Event) -> Self {
        match ev {
            live::Event::ContentReady { hash } => Self::ContentReady { hash },
            live::Event::NeighborUp(peer) => Self::NeighborUp(peer),
            live::Event::NeighborDown(peer) => Self::NeighborDown(peer),
            live::Event::SyncFinished(ev) => Self::SyncFinished(ev),
            live::Event::PendingContentReady => Self::PendingContentReady,
        }
    }
}

impl LiveEvent {
    fn from_replica_event(
        ev: crate::Event,
        content_status_cb: &ContentStatusCallback,
    ) -> Result<Self> {
        Ok(match ev {
            crate::Event::LocalInsert { entry, .. } => Self::InsertLocal {
                entry: entry.into(),
            },
            crate::Event::RemoteInsert { entry, from, .. } => Self::InsertRemote {
                content_status: content_status_cb(entry.content_hash()),
                entry: entry.into(),
                from: PublicKey::from_bytes(&from)?,
            },
        })
    }
}

/// Where to persist the default author.
///
/// If set to `Mem`, a new author will be created in the docs store before spawning the sync
/// engine. Changing the default author will not be persisted.
///
/// If set to `Persistent`, the default author will be loaded from and persisted to the specified
/// path (as hex encoded string of the author's public key).
#[derive(Debug)]
pub enum DefaultAuthorStorage {
    /// Memory storage.
    Mem,
    /// File based persistent storage.
    Persistent(PathBuf),
}

impl DefaultAuthorStorage {
    /// Load the default author from the storage.
    ///
    /// Will create and save a new author if the storage is empty.
    ///
    /// Returns an error if the author can't be parsed or if the uathor does not exist in the docs
    /// store.
    pub(crate) async fn load(&self, docs_store: &SyncHandle) -> anyhow::Result<AuthorId> {
        match self {
            Self::Mem => {
                let author = Author::new(&mut rand::thread_rng());
                let author_id = author.id();
                docs_store.import_author(author).await?;
                Ok(author_id)
            }
            Self::Persistent(ref path) => {
                if path.exists() {
                    let data = tokio::fs::read_to_string(path).await.with_context(|| {
                        format!(
                            "Failed to read the default author file at `{}`",
                            path.to_string_lossy()
                        )
                    })?;
                    let author_id = AuthorId::from_str(&data).with_context(|| {
                        format!(
                            "Failed to parse the default author from `{}`",
                            path.to_string_lossy()
                        )
                    })?;
                    if docs_store.export_author(author_id).await?.is_none() {
                        bail!("The default author is missing from the docs store. To recover, delete the file `{}`. Then iroh will create a new default author.", path.to_string_lossy())
                    }
                    Ok(author_id)
                } else {
                    let author = Author::new(&mut rand::thread_rng());
                    let author_id = author.id();
                    docs_store.import_author(author).await?;
                    // Make sure to write the default author to the store
                    // *before* we write the default author ID file.
                    // Otherwise the default author ID file is effectively a dangling reference.
                    docs_store.flush_store().await?;
                    self.persist(author_id).await?;
                    Ok(author_id)
                }
            }
        }
    }

    /// Save a new default author.
    pub async fn persist(&self, author_id: AuthorId) -> anyhow::Result<()> {
        match self {
            Self::Mem => {
                // persistence is not possible for the mem storage so this is a noop.
            }
            Self::Persistent(ref path) => {
                tokio::fs::write(path, author_id.to_string())
                    .await
                    .with_context(|| {
                        format!(
                            "Failed to write the default author to `{}`",
                            path.to_string_lossy()
                        )
                    })?;
            }
        }
        Ok(())
    }
}

/// Persistent default author for a docs engine.
#[derive(Debug)]
pub struct DefaultAuthor {
    value: RwLock<AuthorId>,
    storage: DefaultAuthorStorage,
}

impl DefaultAuthor {
    /// Load the default author from storage.
    ///
    /// If the storage is empty creates a new author and persists it.
    pub(crate) async fn load(
        storage: DefaultAuthorStorage,
        docs_store: &SyncHandle,
    ) -> Result<Self> {
        let value = storage.load(docs_store).await?;
        Ok(Self {
            value: RwLock::new(value),
            storage,
        })
    }

    /// Get the current default author.
    pub async fn get(&self) -> AuthorId {
        *self.value.read().await
    }

    /// Set the default author.
    pub(crate) async fn set(&self, author_id: AuthorId, docs_store: &SyncHandle) -> Result<()> {
        if docs_store.export_author(author_id).await?.is_none() {
            bail!("The author does not exist");
        }
        self.storage.persist(author_id).await?;
        *self.value.write().await = author_id;
        Ok(())
    }
}