rings-core 0.20.0

use std::str::FromStr;

use async_trait::async_trait;
use rings_transport::core::transport::MAX_DATA_CHANNEL_MESSAGE_SIZE;
use serde::Serialize;

use super::StorageSyncDestination;
use super::StorageSyncPurpose;
use super::StorageSyncTarget;
use crate::consts::MAX_CHUNK_ENVELOPE_OVERHEAD;
use crate::consts::TRANSPORT_CUSTOM_OVERHEAD;
use crate::dht::chord::PeerRing;
use crate::dht::chord::PeerRingAction;
use crate::dht::did::BiasId;
use crate::dht::entry::Entry;
use crate::dht::entry::PlacedEntry;
use crate::dht::entry::SyncedEntryAck;
use crate::dht::ChordStorageSync;
use crate::dht::Did;
use crate::error::Error;
use crate::error::Result;
use crate::message::types::Message;
use crate::message::types::SyncEntriesWithSuccessor;

/// Maximum wire budget for one `SyncEntriesWithSuccessor` hand-off batch.
///
/// This stays below one interoperable WebRTC data-channel frame so storage
/// anti-entropy cannot monopolize the chunk sender. The batch cost also
/// reserves the payload/chunk envelope bytes below.
pub(crate) const SYNC_BATCH_MAX_BYTES: usize = MAX_DATA_CHANNEL_MESSAGE_SIZE / 4;

const SYNC_BATCH_ENVELOPE_HEADROOM_BYTES: usize =
    MAX_CHUNK_ENVELOPE_OVERHEAD + TRANSPORT_CUSTOM_OVERHEAD;

fn serialized_wire_size<T: Serialize>(value: &T) -> Result<usize> {
    let bytes = rings_codec::serialized_size(value).map_err(Error::CodecSerialize)?;
    usize::try_from(bytes).map_err(|_| Error::MessageSizeOverflow)
}

fn add_wire_cost(total: usize, next: usize) -> Result<usize> {
    total.checked_add(next).ok_or(Error::MessageSizeOverflow)
}

fn sync_entries_fixed_wire_cost() -> Result<usize> {
    let empty_message = Message::SyncEntriesWithSuccessor(SyncEntriesWithSuccessor {
        purpose: StorageSyncPurpose::OwnershipHandoff,
        destination: StorageSyncDestination::PhysicalOwner(Did::from(0u32)),
        data: Vec::new(),
    });
    add_wire_cost(
        serialized_wire_size(&empty_message)?,
        SYNC_BATCH_ENVELOPE_HEADROOM_BYTES,
    )
}

fn placed_entry_wire_cost(placed: &PlacedEntry) -> Result<usize> {
    serialized_wire_size(placed)
}

#[cfg(all(test, not(all(feature = "wasm", target_family = "wasm"))))]
pub(super) fn sync_entries_batch_wire_cost(data: &[PlacedEntry]) -> Result<usize> {
    let mut cost = sync_entries_fixed_wire_cost()?;
    for placed in data {
        cost = add_wire_cost(cost, placed_entry_wire_cost(placed)?)?;
    }
    Ok(cost)
}

pub(super) fn sync_entries_batches(
    data: Vec<PlacedEntry>,
    max_batch_bytes: usize,
) -> Result<Vec<Vec<PlacedEntry>>> {
    let mut batches = Vec::new();
    let mut current = Vec::new();
    let fixed_cost = sync_entries_fixed_wire_cost()?;
    let mut current_cost = fixed_cost;

    // Pre: `data` is the migrating set M produced from local storage.
    // Post Coverage: concatenating all returned batches yields exactly M in
    // the same order; no PlacedEntry is duplicated or dropped.
    // Post Budget: every non-singleton batch, and every singleton whose own
    // cost fits, has sync_entries_batch_wire_cost(batch) <= max_batch_bytes.
    // Post Atomicity: each PlacedEntry is moved as a whole; no entry is split
    // across batches.
    // Post Progress: if one PlacedEntry exceeds max_batch_bytes by itself, it
    // is emitted as a one-entry batch so the chunk layer can still frame it.
    for placed in data {
        let placed_cost = placed_entry_wire_cost(&placed)?;
        let candidate_cost = add_wire_cost(current_cost, placed_cost)?;
        if current.is_empty() {
            current.push(placed);
            current_cost = candidate_cost;
            continue;
        }

        if candidate_cost <= max_batch_bytes {
            current.push(placed);
            current_cost = candidate_cost;
        } else {
            batches.push(current);
            current = vec![placed];
            current_cost = add_wire_cost(fixed_cost, placed_cost)?;
        }
    }

    if !current.is_empty() {
        batches.push(current);
    }

    Ok(batches)
}

#[cfg_attr(all(feature = "wasm", target_family = "wasm"), async_trait(?Send))]
#[cfg_attr(not(all(feature = "wasm", target_family = "wasm")), async_trait)]
impl ChordStorageSync<PeerRingAction> for PeerRing {
    /// When the successor of a node is updated, it needs to check if there are
    /// `Entry`s that are no longer between current node and `new_successor`,
    /// and copy them to the new successor.
    async fn sync_entries_with_successor(&self, new_successor: Did) -> Result<PeerRingAction> {
        if self.storage_virtual_nodes_enabled()? {
            return self.copy_entries_to_observed_virtual_storage_owners().await;
        }

        let mut data = Vec::<PlacedEntry>::new();
        let all_items: Vec<(String, Entry)> = self.storage.get_all().await?;

        // Pre: new_successor is the successor adopted by stabilization.
        // Post S1: forall key in local_before, local_after[key] =
        // local_before[key]; this transition emits join deliveries only.
        // Post S2(copy): every emitted PlacedEntry keeps the exact local
        // placement key, so an eventual ack names the key whose durable copy was
        // reported by the receiver.
        // Preservation #611/#614: sync hand-off is join-before-ack-before-local
        // cleanup. acknowledge_synced_entries is the only local cleanup
        // transition and does not define storage convergence.
        for (entry_key_str, entry) in all_items {
            let entry_key = Did::from_str(&entry_key_str)?;
            if BiasId::cmp_from_observer(self.did, entry_key, new_successor)
                == std::cmp::Ordering::Greater
            {
                data.push(PlacedEntry::new(entry_key, entry));
            }
        }

        let batches = sync_entries_batches(data, SYNC_BATCH_MAX_BYTES)?;
        Ok(batches
            .into_iter()
            .map(|batch| {
                PeerRingAction::sync_entries_for_handoff(
                    StorageSyncDestination::PhysicalOwner(new_successor),
                    batch,
                )
            })
            .collect::<Vec<_>>()
            .into())
    }

    async fn acknowledge_synced_entries(&self, acks: &[SyncedEntryAck]) -> Result<PeerRingAction> {
        // Pre S2': each ack in acks is contained in a
        // SyncEntriesWithSuccessorReport sent only after the receiver persisted
        // SyncedEntryAck { key, entry } at key.
        // Post S2': a local key is removed only if canonical(local_before[key])
        // == canonical(ack.entry). If the canonical local value differs, the
        // local value is preserved and will be offered again by a later
        // sync_entries_with_successor transition.
        // Preservation #614: a write racing between copy and ack changes the
        // canonical local value, so confirms_local_value is false and delete
        // is skipped.
        for ack in acks {
            let Some(local_entry) = self.storage.get(&ack.key.to_string()).await? else {
                continue;
            };
            if ack.confirms_local_value(&local_entry)? {
                self.storage.remove(&ack.key.to_string()).await?;
            }
        }

        Ok(PeerRingAction::None)
    }
}

impl PeerRing {
    async fn copy_entries_to_observed_virtual_storage_owners(&self) -> Result<PeerRingAction> {
        let all_items: Vec<(String, Entry)> = self.storage.get_all().await?;
        let mut by_target =
            std::collections::BTreeMap::<StorageSyncDestination, Vec<PlacedEntry>>::new();

        // Pre: storage virtual nodes are enabled.
        // Model: storage_sync_target computes ownership under this node's
        // authenticated local topology view, not a globally complete membership
        // relation.
        // Post S1: local entries are retained without action; entries whose
        // observed virtual owner is remote are emitted as additive anti-entropy
        // copies to that physical owner.
        // Preservation S1'': this transition cannot create a delete-capable
        // report. Only non-virtual physical handoff has an ownership proof
        // strong enough to permit source cleanup.
        for (entry_key_str, entry) in all_items {
            let entry_key = Did::from_str(&entry_key_str)?;
            if let StorageSyncTarget::Remote(target) = self.storage_sync_target(entry_key)? {
                by_target
                    .entry(target)
                    .or_default()
                    .push(PlacedEntry::new(entry_key, entry));
            }
        }

        let mut actions = Vec::new();
        for (target, data) in by_target {
            for batch in sync_entries_batches(data, SYNC_BATCH_MAX_BYTES)? {
                actions.push(PeerRingAction::sync_entries_for_repair(target, batch));
            }
        }
        Ok(actions.into())
    }
}