nodedb-lite 0.1.0

Embedded edge client for NodeDB — offline vector search, graph traversal, and CRDT sync
Documentation
//! Control / latched single-shot messages: reactive token refresh, pending
//! resync requests, row push refusals, pending array acks, array catch-up
//! requests, and the CRDT delta push (the original sync flow). Drained once
//! per tick before the per-engine queues.

use std::ops::ControlFlow;
use std::sync::Arc;

use futures::SinkExt;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;

use nodedb_types::hlc::Hlc;
use nodedb_types::sync::wire::{
    CollectionSchemaSyncMsg, EngineKind, SyncFrame, SyncMessageType, stream_id_for,
};

use super::send::{encode_and_send, send_binary};
use crate::sync::client::SyncClient;
use crate::sync::collection_schema_builder::descriptor_from_meta;
use crate::sync::delegate::SyncDelegate;

/// Drain control messages: token refresh (when paused for auth), resync
/// requests, row push refusals, array acks, and array catch-up requests.
/// Returns `Break` if push must pause this tick
/// (auth pause) or the connection is lost.
pub(super) async fn push_control_messages<S>(
    client: &Arc<SyncClient>,
    sink: &Arc<Mutex<S>>,
) -> ControlFlow<()>
where
    S: SinkExt<Message> + Unpin,
    S::Error: std::fmt::Display,
{
    if client.is_push_paused_for_auth().await {
        // Only attempt a refresh if the backoff interval has elapsed.
        if client.is_refresh_backoff_elapsed().await
            && let Some(refresh_msg) = client.initiate_token_refresh().await
        {
            encode_and_send(
                sink,
                SyncMessageType::TokenRefresh,
                &refresh_msg,
                "TokenRefresh (reactive)",
            )
            .await?;
        }
        // Paused — emit nothing else this tick.
        return ControlFlow::Break(());
    }

    if let Some(resync) = client.take_pending_resync().await {
        encode_and_send(
            sink,
            SyncMessageType::ResyncRequest,
            &resync,
            "ResyncRequest",
        )
        .await?;
        tracing::info!(
            reason = ?resync.reason,
            from_mutation_id = resync.from_mutation_id,
            "sent ResyncRequest to Origin"
        );
    }

    let mut rejects = client.drain_row_push_rejects().await.into_iter();
    while let Some(reject) = rejects.next() {
        let Some(frame) = SyncFrame::try_encode(SyncMessageType::RowPushReject, &reject) else {
            tracing::error!(
                collection = %reject.collection,
                document_id = %reject.document_id,
                "failed to encode RowPushReject frame; Origin is not told about this refused row"
            );
            continue;
        };
        if let Err(e) = send_binary(sink, frame).await {
            tracing::warn!(
                collection = %reject.collection,
                document_id = %reject.document_id,
                error = %e,
                "RowPushReject send failed; the unsent refusals are re-sent after reconnect"
            );
            client
                .requeue_row_push_rejects(std::iter::once(reject).chain(rejects).collect())
                .await;
            return ControlFlow::Break(());
        }
        tracing::debug!(
            collection = %reject.collection,
            document_id = %reject.document_id,
            "sent RowPushReject to Origin"
        );
    }

    for ack in client.drain_pending_array_acks().await {
        let Some(frame) = SyncFrame::try_encode(SyncMessageType::ArrayAck, &ack) else {
            tracing::warn!(array = %ack.array, "failed to encode ArrayAck frame; dropping");
            continue;
        };
        if let Err(e) = send_binary(sink, frame).await {
            tracing::warn!(array = %ack.array, error = %e, "ArrayAck send failed");
            return ControlFlow::Break(());
        }
        tracing::debug!(array = %ack.array, "sent ArrayAck to Origin");
    }

    // A send error drops the unsent requests: the reconnect rebuilds the full
    // round from the durable catch-up tracker, flags included.
    for request in client.drain_array_catchups().await {
        let Some(frame) = SyncFrame::try_encode(SyncMessageType::ArrayCatchupRequest, &request)
        else {
            tracing::error!(
                array = %request.array,
                "failed to encode ArrayCatchupRequest frame; array not caught up this session"
            );
            continue;
        };
        if let Err(e) = send_binary(sink, frame).await {
            tracing::warn!(array = %request.array, error = %e, "ArrayCatchupRequest send failed");
            return ControlFlow::Break(());
        }
        tracing::debug!(array = %request.array, "sent ArrayCatchupRequest to Origin");
    }

    ControlFlow::Continue(())
}

/// Announce one collection's schema (`CollectionSchema`, opcode `0x13`) to
/// Origin if it has not already been announced in this session, so Origin
/// materializes the collection before any of its engine data arrives. Mirrors
/// Origin's announce-before-shape-snapshot ordering in
/// `session_handler/announce.rs`.
///
/// Idempotent per session via the client's `announced_collections` set: the
/// first call for a name sends the frame and records it; subsequent calls
/// short-circuit. Called by every outbound engine push (CRDT, columnar,
/// vector, fts, spatial, timeseries) immediately before it emits data for a
/// collection, so the announce ordering holds regardless of which engine
/// carries the collection's first write. A missing/unreconstructable
/// descriptor is skipped (with a warning from `descriptor_from_meta`) rather
/// than emitting an incorrect schema; a send failure returns `Break` so the
/// caller stops pushing this tick.
pub(super) async fn ensure_collection_announced<S>(
    client: &Arc<SyncClient>,
    delegate: &Arc<dyn SyncDelegate>,
    sink: &Arc<Mutex<S>>,
    name: &str,
) -> ControlFlow<()>
where
    S: SinkExt<Message> + Unpin,
    S::Error: std::fmt::Display,
{
    {
        let announced = client.announced_collections().lock().await;
        if announced.contains(name) {
            return ControlFlow::Continue(());
        }
    }

    let Some(meta) = delegate.get_collection_meta(name).await else {
        tracing::debug!(
            collection = %name,
            "no persisted metadata; skipping schema announce (implicit collection)"
        );
        return ControlFlow::Continue(());
    };
    let Some(descriptor) = descriptor_from_meta(&meta, client.config().tenant_id) else {
        // descriptor_from_meta already warned with the specific reason.
        return ControlFlow::Continue(());
    };

    let msg = CollectionSchemaSyncMsg {
        descriptor,
        creation_hlc: Hlc::ZERO,
    };
    let Some(frame) = SyncFrame::try_encode(SyncMessageType::CollectionSchema, &msg) else {
        tracing::error!(collection = %name, "failed to encode CollectionSchema frame; skipping");
        return ControlFlow::Continue(());
    };
    if let Err(e) = send_binary(sink, frame).await {
        tracing::warn!(collection = %name, error = %e, "CollectionSchema send failed");
        return ControlFlow::Break(());
    }

    client
        .announced_collections()
        .lock()
        .await
        .insert(name.to_string());
    tracing::debug!(collection = %name, "announced CollectionSchema to Origin");

    ControlFlow::Continue(())
}

/// Announce collection schemas for every collection with pending CRDT deltas
/// that hasn't already been announced in this session, so Origin materializes
/// the collection before its deltas arrive.
pub(in crate::sync::transport) async fn push_collection_schemas<S>(
    client: &Arc<SyncClient>,
    delegate: &Arc<dyn SyncDelegate>,
    sink: &Arc<Mutex<S>>,
) -> ControlFlow<()>
where
    S: SinkExt<Message> + Unpin,
    S::Error: std::fmt::Display,
{
    let pending = delegate.pending_deltas();
    if pending.is_empty() {
        return ControlFlow::Continue(());
    }

    let mut names: Vec<String> = pending.into_iter().map(|d| d.collection).collect();
    names.sort_unstable();
    names.dedup();

    for name in names {
        ensure_collection_announced(client, delegate, sink, &name).await?;
    }

    ControlFlow::Continue(())
}

/// Push pending CRDT deltas, respecting the flow control window.
pub(in crate::sync::transport) async fn push_crdt_deltas<S>(
    client: &Arc<SyncClient>,
    delegate: &Arc<dyn SyncDelegate>,
    sink: &Arc<Mutex<S>>,
) -> ControlFlow<()>
where
    S: SinkExt<Message> + Unpin,
    S::Error: std::fmt::Display,
{
    let pending = delegate.pending_deltas();
    if pending.is_empty() {
        return ControlFlow::Continue(());
    }

    let pending_bytes: usize = pending.iter().map(|d| d.delta_bytes.len()).sum();
    client
        .update_pending_stats(pending.len(), pending_bytes)
        .await;

    // Read the peer id now, not at connect: a collision refusal earlier in this
    // session rotates it, and the rows queued by that rotation must go out
    // under the new identity or be refused again.
    let mut msgs = client
        .build_delta_pushes(&pending, delegate.sync_identity().peer_id)
        .await;
    if msgs.is_empty() {
        return ControlFlow::Continue(()); // flow control window full — wait for ACKs
    }

    // Stamp each message with producer identity and a stable per-collection
    // stream seq. The seq is assigned once at first send and stored back on
    // the engine's pending delta so that reconnect re-sends reuse the same
    // seq — Origin deduplicates by seq rather than double-applying.
    let seq_by_mid: std::collections::HashMap<u64, u64> =
        pending.iter().map(|d| (d.mutation_id, d.seq)).collect();
    let producer_id = client.producer_id().await;
    let epoch = client.accepted_epoch().await;
    for msg in &mut msgs {
        let stream_id = stream_id_for(EngineKind::Crdt, &msg.collection);
        let seq = match seq_by_mid.get(&msg.mutation_id).copied() {
            Some(s) if s != 0 => s, // reuse stable seq (reconnect re-send)
            _ => {
                let s = delegate.next_stream_seq(stream_id).await;
                delegate.set_pending_delta_seq(msg.mutation_id, s).await; // persist back to engine
                s
            }
        };
        msg.producer_id = producer_id;
        msg.epoch = epoch;
        msg.seq = seq;
    }

    let mutation_ids: Vec<u64> = msgs.iter().map(|m| m.mutation_id).collect();
    {
        let mut sink_guard = sink.lock().await;
        for msg in &msgs {
            let Some(frame) = SyncFrame::try_encode(SyncMessageType::DeltaPush, msg) else {
                tracing::error!("failed to encode delta push frame; dropping batch");
                return ControlFlow::Break(());
            };
            if let Err(e) = sink_guard
                .send(Message::Binary(frame.to_bytes().into()))
                .await
            {
                tracing::warn!(error = %e, "delta push send failed");
                return ControlFlow::Break(()); // connection lost
            }
        }
    }

    client.record_push(&mutation_ids).await;
    tracing::debug!(count = msgs.len(), "pushed deltas to Origin");
    ControlFlow::Continue(())
}