rings_core/dht/

stabilization.rs

//! Stabilization run daemons to maintain dht.

use std::sync::Arc;

use rings_transport::core::transport::WebrtcConnectionState;

use crate::dht::entry::PlacedEntry;
use crate::dht::successor::SuccessorReader;
use crate::dht::types::ChordStorageRepair;
use crate::dht::types::CorrectChord;
use crate::dht::Chord;
use crate::dht::Did;
use crate::dht::PeerRing;
use crate::dht::PeerRingAction;
use crate::dht::PeerRingRemoteAction;
use crate::error::Error;
use crate::error::Result;
use crate::message::FindSuccessorReportHandler;
use crate::message::FindSuccessorSend;
use crate::message::FindSuccessorThen;
use crate::message::Message;
use crate::message::MessagePayload;
use crate::message::NotifyPredecessorSend;
use crate::message::PayloadSender;
use crate::message::QueryForTopoInfoSend;
use crate::message::SyncEntriesWithSuccessor;
use crate::swarm::transport::SwarmTransport;

/// The stabilization runner.
#[derive(Clone)]
pub struct Stabilizer {
    transport: Arc<SwarmTransport>,
    dht: Arc<PeerRing>,
}

impl Stabilizer {
    /// Create a new stabilization runner.
    pub fn new(transport: Arc<SwarmTransport>) -> Self {
        let dht = transport.dht.clone();
        Self { transport, dht }
    }

    /// Run stabilization once.
    pub async fn stabilize(&self) -> Result<()> {
        tracing::debug!("STABILIZATION notify_predecessor start");
        if let Err(e) = self.notify_predecessor().await {
            tracing::error!("[stabilize] Failed on notify predecessor {:?}", e);
        }
        tracing::debug!("STABILIZATION notify_predecessor end");
        tracing::debug!("STABILIZATION fix_fingers start");
        if let Err(e) = self.fix_fingers().await {
            tracing::error!("[stabilize] Failed on fix_finger {:?}", e);
        }
        tracing::debug!("STABILIZATION fix_fingers end");
        tracing::debug!("STABILIZATION clean_unavailable_connections start");
        if let Err(e) = self.clean_unavailable_connections().await {
            tracing::error!(
                "[stabilize] Failed on clean unavailable connections {:?}",
                e
            );
        }
        tracing::debug!("STABILIZATION clean_unavailable_connections end");
        // Default HMCC/Zave stabilization path. The pure operation is specified
        // as `CorrectStabilize` in tests/default/dht_convergence.rs.
        tracing::debug!("STABILIZATION correct_stabilize start");
        if let Err(e) = self.correct_stabilize().await {
            tracing::error!("[stabilize] Failed on call correct stabilize {:?}", e);
        }
        tracing::debug!("STABILIZATION correct_stabilize end");
        tracing::debug!("STABILIZATION repair_storage start");
        if let Err(e) = self.repair_storage().await {
            tracing::error!("[stabilize] Failed on repair storage {:?}", e);
        }
        tracing::debug!("STABILIZATION repair_storage end");
        Ok(())
    }

    fn collect_storage_repair_actions(
        act: PeerRingAction,
        out: &mut Vec<(Did, Vec<PlacedEntry>)>,
    ) -> Result<()> {
        match act {
            PeerRingAction::None => Ok(()),
            PeerRingAction::RemoteAction(
                destination,
                PeerRingRemoteAction::SyncEntriesWithSuccessor(data),
            ) => {
                out.push((destination, data));
                Ok(())
            }
            PeerRingAction::MultiActions(actions) => {
                for action in actions {
                    Self::collect_storage_repair_actions(action, out)?;
                }
                Ok(())
            }
            action => {
                tracing::error!("Invalid storage repair action: {action:?}");
                Err(crate::error::Error::unexpected_peer_ring_action(action))
            }
        }
    }

    async fn handle_storage_repair_action(&self, act: PeerRingAction) -> Result<()> {
        let mut messages = Vec::new();
        Self::collect_storage_repair_actions(act, &mut messages)?;
        for (destination, data) in messages {
            self.transport
                .send_message(
                    Message::SyncEntriesWithSuccessor(SyncEntriesWithSuccessor { data }),
                    destination,
                )
                .await?;
        }
        Ok(())
    }

    /// Republish locally-held entries to their current affine owners.
    pub async fn repair_storage(&self) -> Result<()> {
        let action = self
            .dht
            .republish_local_entries(self.transport.storage_redundancy())
            .await?;
        self.handle_storage_repair_action(action).await
    }

    /// Clean unavailable connections in transport.
    pub async fn clean_unavailable_connections(&self) -> Result<()> {
        let conns = self.transport.get_connections();

        for (did, conn) in conns.into_iter() {
            // Only terminal states are cleaned. `Disconnected` is transient: ICE
            // can recover from it, so tearing it down here (the stabilizer runs
            // every few seconds) would kill connections during a brief blip
            // before WebRTC self-heals. This mirrors the swarm callback, which
            // also only leaves the DHT on `Failed`/`Closed`.
            if matches!(
                conn.webrtc_connection_state(),
                WebrtcConnectionState::Failed | WebrtcConnectionState::Closed
            ) {
                tracing::info!("STABILIZATION clean_unavailable_transports: {:?}", did);
                let should_repair = self
                    .dht
                    .peer_may_share_storage_responsibility(did, self.transport.storage_redundancy())
                    .await?;
                self.transport.disconnect(did).await?;
                if should_repair {
                    self.repair_storage().await?;
                }
            }
        }

        Ok(())
    }

    /// Notify predecessor, this is a DHT operation.
    pub async fn notify_predecessor(&self) -> Result<()> {
        let (successor_min, successor_list) = {
            let successor = self.dht.successors();
            (successor.min()?, successor.list()?)
        };

        let msg = Message::NotifyPredecessorSend(NotifyPredecessorSend { did: self.dht.did });
        if self.dht.did != successor_min {
            for s in successor_list {
                tracing::debug!("STABILIZATION notify_predecessor: {:?}", s);
                let payload =
                    MessagePayload::new_send(msg.clone(), self.transport.session_sk(), s, s)?;
                self.transport.send_payload(payload).await?;
            }
            Ok(())
        } else {
            Ok(())
        }
    }

    /// Fix fingers from finger table, this is a DHT operation.
    async fn fix_fingers(&self) -> Result<()> {
        match self.dht.fix_fingers() {
            Ok(action) => match action {
                PeerRingAction::None => Ok(()),
                PeerRingAction::RemoteAction(
                    closest_predecessor,
                    PeerRingRemoteAction::FindSuccessorForFix {
                        did: finger_did,
                        index,
                    },
                ) => {
                    tracing::debug!("STABILIZATION fix_fingers: {:?}", finger_did);
                    let msg = Message::FindSuccessorSend(FindSuccessorSend {
                        did: finger_did,
                        then: FindSuccessorThen::Report(
                            FindSuccessorReportHandler::FixFingerTable { index },
                        ),
                        strict: false,
                    });
                    let payload = MessagePayload::new_send(
                        msg.clone(),
                        self.transport.session_sk(),
                        closest_predecessor,
                        closest_predecessor,
                    )?;
                    self.transport.send_payload(payload).await?;
                    Ok(())
                }
                _ => {
                    tracing::error!("Invalid PeerRing Action");
                    Err(Error::PeerRingInvalidAction)
                }
            },
            Err(e) => {
                tracing::error!("{:?}", e);
                Err(e)
            }
        }
    }

    /// Call stabilization from correct chord implementation
    pub async fn correct_stabilize(&self) -> Result<()> {
        if let PeerRingAction::RemoteAction(
            next,
            PeerRingRemoteAction::QueryForSuccessorListAndPred,
        ) = self.dht.pre_stabilize()?
        {
            self.transport
                .send_direct_message(
                    Message::QueryForTopoInfoSend(QueryForTopoInfoSend::new_for_stab(next)),
                    next,
                )
                .await?;
        }
        Ok(())
    }
}

#[cfg(not(feature = "wasm"))]
mod stabilizer {
    use std::sync::Arc;
    use std::time::Duration;

    use futures::future::FutureExt;
    use futures::pin_mut;
    use futures::select;
    use futures_timer::Delay;

    use super::*;

    impl Stabilizer {
        /// Run stabilization in a loop.
        pub async fn wait(self: Arc<Self>, interval: Duration) {
            loop {
                let timeout = Delay::new(interval).fuse();
                pin_mut!(timeout);
                select! {
                    _ = timeout => self
                        .stabilize()
                        .await
                        .unwrap_or_else(|e| tracing::error!("failed to stabilize {:?}", e)),
                }
            }
        }
    }
}

#[cfg(feature = "wasm")]
mod stabilizer {
    use std::sync::Arc;
    use std::time::Duration;

    use super::*;
    use crate::poll;

    impl Stabilizer {
        /// Run stabilization in a loop.
        pub async fn wait(self: Arc<Self>, interval: Duration) {
            let millis = i32::try_from(interval.as_millis()).unwrap_or(i32::MAX);
            let stabilizer = self;
            poll!(
                {
                    let stabilizer = Arc::clone(&stabilizer);
                    async move {
                        stabilizer
                            .stabilize()
                            .await
                            .unwrap_or_else(|e| tracing::error!("failed to stabilize {:?}", e));
                    }
                },
                millis
            );
        }
    }
}