nodedb-lite 0.1.0

//! WebSocket connect + handshake. One attempt per `connect_and_run` call;
//! retries are handled by the outer `run_sync_loop` with exponential backoff.

use std::sync::Arc;
use std::time::Duration;

use futures::{SinkExt, StreamExt};
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;

use nodedb_types::sync::wire::{SyncFrame, SyncMessageType};

use super::dispatch::receive_loop;
use super::push::{delta_push_loop, ping_loop};
use crate::error::LiteError;
use crate::sync::client::{SyncClient, SyncState};
use crate::sync::delegate::SyncDelegate;

/// Single connection attempt: connect → handshake → message loop.
///
/// Returns `Ok(())` on a clean server-initiated close and `Err` for any
/// transport, handshake, or read error. Background push and ping tasks are
/// always cancelled before this function returns.
pub(super) async fn connect_and_run(
    client: &Arc<SyncClient>,
    delegate: &Arc<dyn SyncDelegate>,
) -> Result<(), LiteError> {
    // Reload durable producer identity so the outbound path has a valid
    // producer_id/accepted_epoch before the handshake is built. This is a
    // no-op on the very first connect (returns 0, 0) and a restore on
    // reconnect (returns the last values Origin assigned).
    let (persisted_producer_id, persisted_accepted_epoch) = delegate.load_producer_state().await;
    client
        .load_producer_state(persisted_producer_id, persisted_accepted_epoch)
        .await;

    // Reset per-connection inbound state (shape LSN gaps, flow control).
    // The per-stream seq frontier (StreamSeqTracker) is NOT reset here —
    // it is loaded once from storage at startup and never cleared on reconnect
    // so outbound frame numbering resumes from the durable last_assigned.
    // The fenced flag is cleared so this attempt can push; if Origin still
    // fences the producer the flag will be set again on the first ack.
    client.reset_sequence_tracking().await;
    client.reset_flow_control().await;
    client.clear_fenced();

    // Reset the array push window: the new session announces every array's
    // schema again and re-sends every unconfirmed op from the queue head.
    // Ops held for the last session become sendable again.
    client.array_push().lock().await.reset();

    // Clear in-flight maps for all engine outbound queues. The durable entries
    // are still in storage; clearing in-flight makes them eligible for
    // re-drain → re-send → Origin deduplicates via its idempotent gate.
    delegate.clear_engine_in_flight().await;

    // ── Connect ──
    let (ws_stream, _response) = tokio_tungstenite::connect_async(&client.config().url)
        .await
        .map_err(|e| LiteError::Sync {
            detail: format!("WebSocket connect failed: {e}"),
        })?;

    let (mut sink, mut stream) = ws_stream.split();

    // ── Handshake ──
    let handshake = client.build_handshake(&delegate.sync_identity()).await;
    let frame = SyncFrame::try_encode(SyncMessageType::Handshake, &handshake).ok_or_else(|| {
        LiteError::Sync {
            detail: "failed to encode handshake frame".to_string(),
        }
    })?;
    sink.send(Message::Binary(frame.to_bytes().into()))
        .await
        .map_err(|e| LiteError::Sync {
            detail: format!("handshake send failed: {e}"),
        })?;

    let ack_msg = tokio::time::timeout(Duration::from_secs(10), stream.next())
        .await
        .map_err(|_| LiteError::Sync {
            detail: "handshake timeout".to_string(),
        })?
        .ok_or_else(|| LiteError::Sync {
            detail: "connection closed before handshake ack".to_string(),
        })?
        .map_err(|e| LiteError::Sync {
            detail: format!("handshake read error: {e}"),
        })?;

    let ack_bytes = match &ack_msg {
        Message::Binary(b) => b.as_ref(),
        Message::Text(_)
        | Message::Ping(_)
        | Message::Pong(_)
        | Message::Close(_)
        | Message::Frame(_) => {
            return Err(LiteError::Sync {
                detail: "expected binary handshake ack".to_string(),
            });
        }
    };

    let ack_frame = SyncFrame::from_bytes(ack_bytes).ok_or_else(|| LiteError::Sync {
        detail: "invalid handshake ack frame".to_string(),
    })?;

    if ack_frame.msg_type != SyncMessageType::HandshakeAck {
        return Err(LiteError::Sync {
            detail: format!("expected HandshakeAck, got {:?}", ack_frame.msg_type),
        });
    }

    let ack: nodedb_types::sync::wire::HandshakeAckMsg =
        ack_frame.decode_body().ok_or_else(|| LiteError::Sync {
            detail: "failed to decode HandshakeAck".to_string(),
        })?;

    if !client.handle_handshake_ack(&ack).await {
        return Err(LiteError::Sync {
            detail: format!("handshake rejected: {}", ack.error.unwrap_or_default()),
        });
    }

    // Origin holds a history under this `lite_id` that this instance's local
    // state diverged from — a restored backup, a rolled-back device, a cloned
    // store. Resuming the producer stream would push operations Origin cannot
    // order against what it already has, so the instance becomes a new
    // producer and re-announces its state under it. This must happen before
    // any delta is pushed, which is why it sits on the connect path rather
    // than being left to the application to notice in a log.
    if ack.fork_detected {
        tracing::warn!(
            session = %ack.session_id,
            "Origin reports this instance as forked — regenerating the producer identity"
        );
        delegate.regenerate_identity().await;
        // The identity that just changed is the one the accepted handshake was
        // built from, so this session is over: reconnecting presents the new
        // identity from the first frame.
        return Err(LiteError::Sync {
            detail: "producer identity regenerated after fork; reconnecting".to_string(),
        });
    }

    // Durably persist the server-assigned producer identity so it survives
    // process restart and is available on the next reconnect.
    delegate
        .persist_producer_state(ack.producer_id, ack.accepted_epoch)
        .await;

    // Origin keeps no per-replica array cursor across sessions: each session
    // asks for every array it holds from its catch-up high-water mark, and
    // Origin answers with the missed ops (or a snapshot first when the mark
    // is below its GC boundary). The push loop sends these on its first tick,
    // under the same auth-pause gate as every other control frame.
    client
        .reset_array_catchups(delegate.begin_array_catchup())
        .await;

    // ── Message loop ──
    let sink = Arc::new(Mutex::new(sink));

    let push_sink = Arc::clone(&sink);
    let push_client = Arc::clone(client);
    let push_delegate = Arc::clone(delegate);
    let push_handle = tokio::spawn(async move {
        delta_push_loop(&push_client, &push_delegate, &push_sink).await;
    });

    let ping_sink = Arc::clone(&sink);
    let ping_client = Arc::clone(client);
    let ping_handle = tokio::spawn(async move {
        ping_loop(&ping_client, &ping_sink).await;
    });

    // Own the child handles in a drop guard: cancellation of this future
    // (abort from the FFI stop/close path) drops it and aborts both children.
    let children = AbortChildrenOnDrop {
        handles: vec![push_handle, ping_handle],
    };

    let recv_result = receive_loop(client, delegate, &mut stream).await;

    // Normal path: abort children before returning. The guard would do it
    // anyway on drop; this makes shutdown explicit and ordered.
    drop(children);

    client.set_state(SyncState::Disconnected).await;
    recv_result
}

/// Aborts every child task when dropped — including when the parent future is
/// cancelled mid-await, which skips the normal cleanup path above.
///
/// Without this, a cancelled `connect_and_run` detaches `delta_push_loop` and
/// `ping_loop` (their `JoinHandle`s drop), and the children keep polling the
/// sink and database after the caller considers the sync task quiescent.
struct AbortChildrenOnDrop {
    handles: Vec<tokio::task::JoinHandle<()>>,
}

impl Drop for AbortChildrenOnDrop {
    fn drop(&mut self) {
        for handle in &self.handles {
            handle.abort();
        }
    }
}