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;
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 {
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?;
}
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");
}
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(())
}
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 {
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(())
}
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(())
}
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;
let mut msgs = client
.build_delta_pushes(&pending, delegate.sync_identity().peer_id)
.await;
if msgs.is_empty() {
return ControlFlow::Continue(()); }
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, _ => {
let s = delegate.next_stream_seq(stream_id).await;
delegate.set_pending_delta_seq(msg.mutation_id, s).await; 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(()); }
}
}
client.record_push(&mutation_ids).await;
tracing::debug!(count = msgs.len(), "pushed deltas to Origin");
ControlFlow::Continue(())
}