volt_client_grpc/

sync_provider.rs

//! Sync Provider for collaborative document editing with Y.js/yrs CRDT
//!
//! This module provides a SyncProvider that manages bidirectional synchronization
//! of Y.js/yrs documents with a Volt server. It handles:
//! - Initial document synchronization
//! - Real-time updates from other clients
//! - Chunked payload handling for large updates
//! - State machine management (SYNCING → DONE → UPDATE → AWARENESS)
//! - Subdocument synchronization events

use std::sync::Arc;
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_stream::StreamExt;
use yrs::updates::decoder::Decode;
use yrs::updates::encoder::Encode;
use yrs::{Doc, ReadTxn, StateVector, Subscription, Transact, Update, Uuid};

use crate::error::{Result, VoltError};
use crate::proto::volt::{self, SyncState};
use crate::VoltClient;

const MAX_CHUNK_SIZE: usize = 1024 * 1024; // 1MB chunks

/// Origin marker for updates received from the server
/// Used to distinguish server updates from local changes
const SERVER_ORIGIN: &str = "volt-sync-provider";

/// Events emitted by the SyncProvider
#[derive(Debug, Clone)]
pub enum SyncEvent {
    /// Synchronization started
    Syncing,
    /// Initial synchronization complete
    Synced,
    /// Update received and applied
    Updated,
    /// Connection lost
    Disconnected,
    /// Error occurred
    Error(String),
    /// A subdocument was loaded and needs its own sync provider
    SubdocLoaded { guid: Uuid, doc: Doc },
    /// A subdocument was removed
    SubdocRemoved { guid: Uuid },
}

/// Sync Provider state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderState {
    /// Not connected
    Disconnected,
    /// Synchronizing initial state
    Syncing,
    /// Synced and ready for updates
    Synced,
}

/// SyncProvider manages document synchronization with Volt server
pub struct SyncProvider {
    /// The Y.js document (Doc is internally reference-counted)
    doc: Doc,
    /// Database ID
    database_id: String,
    /// Document ID
    document_id: String,
    /// Current provider state
    state: Arc<RwLock<ProviderState>>,
    /// Event sender
    event_tx: mpsc::Sender<SyncEvent>,
    /// Request sender to Volt server
    request_tx: Option<mpsc::Sender<volt::SyncDocumentRequest>>,
    /// Buffer for chunked payloads
    chunk_buffer: Arc<Mutex<Vec<u8>>>,
    /// Read-only mode
    is_read_only: Arc<RwLock<bool>>,
    /// Document update observer subscription (keeps observer alive)
    #[allow(dead_code)]
    update_subscription: Option<Subscription>,
    /// Subdocs observer subscription (keeps observer alive)
    #[allow(dead_code)]
    subdocs_subscription: Option<Subscription>,
    /// Channel for local updates that need to be sent to server
    local_update_tx: Option<mpsc::Sender<Vec<u8>>>,
}

impl SyncProvider {
    /// Create a new SyncProvider
    ///
    /// # Arguments
    /// * `doc` - The Y.js document to sync
    /// * `database_id` - ID of the database containing the document
    /// * `document_id` - ID of the document to sync
    ///
    /// # Returns
    /// A tuple of (SyncProvider, event_receiver)
    pub fn new(
        doc: Doc,
        database_id: impl Into<String>,
        document_id: impl Into<String>,
    ) -> (Self, mpsc::Receiver<SyncEvent>) {
        let (event_tx, event_rx) = mpsc::channel(100);

        let provider = Self {
            doc,
            database_id: database_id.into(),
            document_id: document_id.into(),
            state: Arc::new(RwLock::new(ProviderState::Disconnected)),
            event_tx,
            request_tx: None,
            chunk_buffer: Arc::new(Mutex::new(Vec::new())),
            is_read_only: Arc::new(RwLock::new(false)),
            update_subscription: None,
            subdocs_subscription: None,
            local_update_tx: None,
        };

        (provider, event_rx)
    }

    /// Start synchronization with the Volt server
    ///
    /// # Arguments
    /// * `client` - The VoltClient to use for synchronization
    /// * `read_only` - If true, only receive updates (no sending)
    /// * `read_only_fallback` - If true, downgrade to read-only if write access denied
    pub async fn start(
        &mut self,
        client: &VoltClient,
        read_only: bool,
        read_only_fallback: bool,
    ) -> Result<()> {
        // Set initial state
        *self.state.write().await = ProviderState::Syncing;
        self.event_tx
            .send(SyncEvent::Syncing)
            .await
            .map_err(|_| VoltError::ConnectionError("Event channel closed".into()))?;

        // Get the current state vector from our document (v1 encoding like JavaScript)
        let state_vector = {
            let txn = self.doc.transact();
            txn.state_vector().encode_v1()
        };

        tracing::debug!("Local state vector: {} bytes", state_vector.len());

        // Create the initial sync request
        let sync_start = volt::SyncDocumentRequest {
            payload: Some(volt::sync_document_request::Payload::SyncStart(
                volt::SyncDocumentStart {
                    database_id: self.database_id.clone(),
                    document_id: self.document_id.clone(),
                    state_vector,
                    read_only,
                    read_only_fallback,
                },
            )),
            sync_state: SyncState::Syncing as i32,
        };

        tracing::debug!(
            "Sending sync start request to database: {}, document: {}",
            self.database_id,
            self.document_id
        );

        // Create bidirectional stream with initial request
        let (tx, mut rx) = client.sync_document(sync_start).await?;
        self.request_tx = Some(tx.clone());

        tracing::debug!("Sync stream established, waiting for responses...");

        // Set up local update observer to send changes to server
        // Create a channel for local updates (observer is sync, sending is async)
        let (local_update_tx, mut local_update_rx) = mpsc::channel::<Vec<u8>>(100);
        self.local_update_tx = Some(local_update_tx.clone());

        // Set up the document observer
        let update_state = self.state.clone();
        let subscription = self
            .doc
            .observe_update_v2(move |txn, event| {
                // Check if this update came from the server (has our origin marker)
                if let Some(origin) = txn.origin() {
                    if origin.as_ref() == SERVER_ORIGIN.as_bytes() {
                        tracing::debug!("Skipping server-originated update");
                        return;
                    }
                }

                // Only send if we're synced (not during initial sync or disconnected)
                // Note: We can't await here, so we use try_send
                let state_guard = update_state.try_read();
                if let Ok(state) = state_guard {
                    if *state == ProviderState::Synced {
                        let update = event.update.clone();
                        tracing::debug!(
                            "Local document updated, queueing {} bytes to send",
                            update.len()
                        );
                        if let Err(e) = local_update_tx.try_send(update) {
                            tracing::error!("Failed to queue local update: {}", e);
                        }
                    } else {
                        tracing::debug!(
                            "Ignoring local update, not synced yet (state: {:?})",
                            *state
                        );
                    }
                }
            })
            .map_err(|e| {
                VoltError::ConnectionError(format!("Failed to set up document observer: {:?}", e))
            })?;

        self.update_subscription = Some(subscription);

        // Set up subdocs observer to emit events when subdocs are loaded/removed
        // The user is responsible for creating SyncProviders for subdocs
        let subdoc_event_tx = self.event_tx.clone();
        let subdocs_sub = self
            .doc
            .observe_subdocs(move |_txn, event| {
                // Emit events for loaded subdocs
                for subdoc in event.loaded() {
                    let guid = subdoc.guid();
                    let subdoc_clone = subdoc.clone();
                    tracing::debug!("Subdoc loaded: {}", guid);
                    if let Err(e) = subdoc_event_tx.try_send(SyncEvent::SubdocLoaded {
                        guid,
                        doc: subdoc_clone,
                    }) {
                        tracing::error!("Failed to send subdoc loaded event: {}", e);
                    }
                }

                // Emit events for removed subdocs
                for subdoc in event.removed() {
                    let guid = subdoc.guid();
                    tracing::debug!("Subdoc removed: {}", guid);
                    if let Err(e) = subdoc_event_tx.try_send(SyncEvent::SubdocRemoved { guid }) {
                        tracing::error!("Failed to send subdoc removed event: {}", e);
                    }
                }
            })
            .map_err(|e| {
                VoltError::ConnectionError(format!("Failed to set up subdocs observer: {:?}", e))
            })?;

        self.subdocs_subscription = Some(subdocs_sub);

        // Spawn task to send local updates to server
        let update_request_tx = tx.clone();
        let update_is_read_only = self.is_read_only.clone();
        tokio::spawn(async move {
            while let Some(update) = local_update_rx.recv().await {
                // Check if read-only
                if *update_is_read_only.read().await {
                    tracing::debug!("Ignoring local update, document is read-only");
                    continue;
                }

                // Send the update
                tracing::debug!("Sending local update to server: {} bytes", update.len());
                if let Err(e) = Self::send_update_chunks(&update_request_tx, &update).await {
                    tracing::error!("Failed to send local update: {}", e);
                }
            }
            tracing::debug!("Local update sender task finished");
        });

        // Spawn task to handle responses
        let doc = self.doc.clone();
        let event_tx = self.event_tx.clone();
        let state = self.state.clone();
        let chunk_buffer = self.chunk_buffer.clone();
        let is_read_only = self.is_read_only.clone();
        let request_tx = tx.clone();

        tokio::spawn(async move {
            tracing::debug!("Response handler task started");
            while let Some(result) = rx.next().await {
                tracing::debug!("Received response from server");
                match result {
                    Ok(response) => {
                        tracing::debug!("Response sync_state: {}", response.sync_state);
                        if let Err(e) = Self::handle_response(
                            &doc,
                            response,
                            &event_tx,
                            &state,
                            &chunk_buffer,
                            &is_read_only,
                            &request_tx,
                        )
                        .await
                        {
                            let _ = event_tx.send(SyncEvent::Error(e.to_string())).await;
                        }
                    }
                    Err(e) => {
                        let _ = event_tx
                            .send(SyncEvent::Error(format!("Stream error: {}", e)))
                            .await;
                        *state.write().await = ProviderState::Disconnected;
                        let _ = event_tx.send(SyncEvent::Disconnected).await;
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// Handle a response from the server
    async fn handle_response(
        doc: &Doc,
        response: volt::SyncDocumentResponse,
        event_tx: &mpsc::Sender<SyncEvent>,
        state: &Arc<RwLock<ProviderState>>,
        chunk_buffer: &Arc<Mutex<Vec<u8>>>,
        is_read_only: &Arc<RwLock<bool>>,
        request_tx: &mpsc::Sender<volt::SyncDocumentRequest>,
    ) -> Result<()> {
        // Check for errors
        if let Some(status) = response.status {
            if status.code != 0 {
                return Err(VoltError::ServerError(format!(
                    "Server error: {}",
                    status.message
                )));
            }
        }

        // Handle read-only mode (only in initial response)
        if response.is_read_only {
            *is_read_only.write().await = true;
        }

        // Get sync state from response
        let sync_state = SyncState::try_from(response.sync_state).unwrap_or(SyncState::Syncing);

        match sync_state {
            SyncState::Syncing => {
                // Server response to our initial sync request
                // Contains state_vector and any updates since last sync
                tracing::debug!("Received SYNC_STATE_SYNCING response");

                // Buffer the update chunk
                if let Some(update) = &response.update {
                    let mut buffer = chunk_buffer.lock().await;
                    buffer.extend_from_slice(&update.chunk);

                    if update.complete {
                        // Apply the server's updates to our document (V2 encoding)
                        // Use SERVER_ORIGIN so the observer knows to skip these updates
                        if !buffer.is_empty() {
                            tracing::debug!(
                                "Applying {} bytes of updates from server",
                                buffer.len()
                            );
                            match Update::decode_v2(&buffer) {
                                Ok(update) => {
                                    let mut txn = doc.transact_mut_with(SERVER_ORIGIN);
                                    if let Err(e) = txn.apply_update(update) {
                                        tracing::error!("Failed to apply update: {:?}", e);
                                    }
                                }
                                Err(e) => {
                                    tracing::error!("Failed to decode update: {:?}", e);
                                }
                            }
                        }
                        buffer.clear();

                        // Now send our local updates based on server's state vector
                        // The server's state_vector tells us what it already has
                        let server_state_vector = if !response.state_vector.is_empty() {
                            match StateVector::decode_v1(&response.state_vector) {
                                Ok(sv) => sv,
                                Err(e) => {
                                    tracing::error!(
                                        "Failed to decode server state vector: {:?}",
                                        e
                                    );
                                    StateVector::default()
                                }
                            }
                        } else {
                            StateVector::default()
                        };

                        // Encode our state as an update relative to the server's state vector (V2)
                        let local_update = {
                            let txn = doc.transact();
                            txn.encode_state_as_update_v2(&server_state_vector)
                        };

                        tracing::debug!(
                            "Sending {} bytes of local updates to server",
                            local_update.len()
                        );

                        // Send DONE message with our local updates
                        let done_request = volt::SyncDocumentRequest {
                            payload: Some(volt::sync_document_request::Payload::Update(
                                volt::SyncDocumentUpdate {
                                    chunk: local_update,
                                    complete: true,
                                },
                            )),
                            sync_state: SyncState::Done as i32,
                        };

                        request_tx.send(done_request).await.map_err(|_| {
                            VoltError::ConnectionError("Failed to send DONE message".into())
                        })?;
                    }
                } else {
                    // No update in response, still need to send DONE with our state
                    let server_state_vector = if !response.state_vector.is_empty() {
                        match StateVector::decode_v1(&response.state_vector) {
                            Ok(sv) => sv,
                            Err(e) => {
                                tracing::error!("Failed to decode server state vector: {:?}", e);
                                StateVector::default()
                            }
                        }
                    } else {
                        StateVector::default()
                    };

                    let local_update = {
                        let txn = doc.transact();
                        txn.encode_state_as_update_v2(&server_state_vector)
                    };

                    tracing::debug!(
                        "Sending {} bytes of local updates (no server update)",
                        local_update.len()
                    );

                    let done_request = volt::SyncDocumentRequest {
                        payload: Some(volt::sync_document_request::Payload::Update(
                            volt::SyncDocumentUpdate {
                                chunk: local_update,
                                complete: true,
                            },
                        )),
                        sync_state: SyncState::Done as i32,
                    };

                    request_tx.send(done_request).await.map_err(|_| {
                        VoltError::ConnectionError("Failed to send DONE message".into())
                    })?;
                }
            }
            SyncState::Done => {
                // Server acknowledges sync complete - we're now synced!
                tracing::debug!("Received SYNC_STATE_DONE response");
                *state.write().await = ProviderState::Synced;
                let _ = event_tx.send(SyncEvent::Synced).await;
            }
            SyncState::Update => {
                // Server sent an update - apply it to our document
                // Use SERVER_ORIGIN so the observer knows to skip these updates
                if let Some(update) = &response.update {
                    let mut buffer = chunk_buffer.lock().await;
                    buffer.extend_from_slice(&update.chunk);

                    if update.complete {
                        if !buffer.is_empty() {
                            tracing::debug!(
                                "Applying {} bytes of live update from server",
                                buffer.len()
                            );
                            match Update::decode_v2(&buffer) {
                                Ok(update) => {
                                    let mut txn = doc.transact_mut_with(SERVER_ORIGIN);
                                    if let Err(e) = txn.apply_update(update) {
                                        tracing::error!("Failed to apply live update: {:?}", e);
                                    }
                                }
                                Err(e) => {
                                    tracing::error!("Failed to decode live update: {:?}", e);
                                }
                            }
                        }
                        let _ = event_tx.send(SyncEvent::Updated).await;
                        buffer.clear();
                    }
                }
            }
            SyncState::Awareness => {
                // Awareness update - ignore for now
                tracing::debug!("Received SYNC_STATE_AWARENESS response");
            }
            _ => {
                tracing::warn!("Unknown sync state: {:?}", sync_state);
            }
        }

        Ok(())
    }

    /// Send an update to the server
    ///
    /// This should be called when the local document changes
    pub async fn send_update(&self, update: &[u8]) -> Result<()> {
        let tx = self
            .request_tx
            .as_ref()
            .ok_or_else(|| VoltError::ConnectionError("Not connected".into()))?;

        // Check if we're read-only
        if *self.is_read_only.read().await {
            return Err(VoltError::PermissionDenied("Document is read-only".into()));
        }

        Self::send_update_chunks(tx, update).await
    }

    /// Send update chunks to the server (static helper for spawned tasks)
    async fn send_update_chunks(
        tx: &mpsc::Sender<volt::SyncDocumentRequest>,
        update: &[u8],
    ) -> Result<()> {
        // Split into chunks if needed
        let chunks = Self::chunk_update_static(update);

        for (i, chunk) in chunks.iter().enumerate() {
            let is_last = i == chunks.len() - 1;

            let request = volt::SyncDocumentRequest {
                payload: Some(volt::sync_document_request::Payload::Update(
                    volt::SyncDocumentUpdate {
                        chunk: chunk.clone(),
                        complete: is_last,
                    },
                )),
                sync_state: SyncState::Update as i32,
            };

            tx.send(request)
                .await
                .map_err(|_| VoltError::ConnectionError("Failed to send update".into()))?;
        }

        Ok(())
    }

    /// Split an update into chunks (static version)
    fn chunk_update_static(update: &[u8]) -> Vec<Vec<u8>> {
        if update.len() <= MAX_CHUNK_SIZE {
            return vec![update.to_vec()];
        }

        update
            .chunks(MAX_CHUNK_SIZE)
            .map(|chunk| chunk.to_vec())
            .collect()
    }

    /// Get the current provider state
    pub async fn state(&self) -> ProviderState {
        *self.state.read().await
    }

    /// Check if the provider is read-only
    pub async fn is_read_only(&self) -> bool {
        *self.is_read_only.read().await
    }

    /// Stop synchronization
    pub async fn stop(&mut self) {
        *self.state.write().await = ProviderState::Disconnected;
        self.request_tx = None;
        self.local_update_tx = None;
        self.update_subscription = None;
        let _ = self.event_tx.send(SyncEvent::Disconnected).await;
    }
}