Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178//! Keeping the connection to the node alive.//!//! The server talks to one node, named on the command line, and owns that//! connection for its whole life. A dev node gets restarted constantly, so the//! link dropping is routine rather than exceptional: a background loop retries//! with backoff, and a tool call about to use the connection retries//! immediately rather than waiting out the loop's timer.
use crate::connection::ConnectionManager;use crate::error::ConnectionResult;use crate::logs;use crate::rpc::{self, rpc_call};use std::sync::Arc;use std::sync::atomic::{AtomicU64, Ordering};use std::time::Duration;use tokio::sync::{Mutex, Notify};use tokio::task::JoinHandle;
/// How long to wait before the first retry after a failed attempt.const MIN_BACKOFF: Duration = Duration::from_millis(250);
/// The longest the background loop will wait between retries.const MAX_BACKOFF: Duration = Duration::from_secs(5);
/// How often the loop re-checks a connection that is up.const HEALTHY_INTERVAL: Duration = Duration::from_secs(2);
/// What the node logs when we reattach to it, so that a restart is visible in/// `get_logs` rather than showing up as a silently emptied buffer.const RECONNECT_NOTICE: &str = "beamdev reconnected: any log events, loaded modules, \ or process state from before this point are gone";
/// Doubles `backoff`, up to [`MAX_BACKOFF`].fn next_backoff(backoff: Duration) -> Duration { (backoff * 2).min(MAX_BACKOFF)}
/// Keeps the manager's connection up, re-establishing it when it drops.#[derive(Debug)]pub struct Supervisor { cookie: String, manager: Arc<ConnectionManager>, /// Held for the duration of an attempt, so the loop and a tool call racing /// to heal the same drop produce one handshake rather than two. attempt: Mutex<()>, /// Wakes the loop early, so a tool call that finds the node down does not /// leave the loop sleeping out a backoff it no longer needs. wake: Notify, /// How many times we have connected, counting the first. connects: AtomicU64,}
impl Supervisor { /// Creates a supervisor for the manager's node, not yet connected. pub fn new(cookie: String, manager: Arc<ConnectionManager>) -> Arc<Self> { Arc::new(Self { cookie, manager, attempt: Mutex::new(()), wake: Notify::new(), connects: AtomicU64::new(0), }) }
/// The node this server is bound to. pub(crate) fn node(&self) -> &str { self.manager.node() }
/// How many times we have connected, counting the first. pub fn connects(&self) -> u64 { self.connects.load(Ordering::SeqCst) }
/// Connects if the link is down, and returns once it is up. /// /// Every tool calls this before using the connection. pub async fn ensure_connected(&self) -> ConnectionResult<()> { if self.manager.is_connected().await { return Ok(()); }
self.wake.notify_one(); self.connect().await }
/// Starts the background loop that heals a drop nobody has asked about yet. pub fn spawn(self: &Arc<Self>) -> JoinHandle<()> { let supervisor = Arc::clone(self); tokio::spawn(async move { supervisor.run().await }) }
async fn run(&self) { let mut backoff = MIN_BACKOFF;
loop { let delay = if self.manager.is_connected().await { backoff = MIN_BACKOFF; HEALTHY_INTERVAL } else { match self.connect().await { Ok(()) => { backoff = MIN_BACKOFF; HEALTHY_INTERVAL } Err(e) => { tracing::debug!("Cannot reach {}: {e}", self.node()); let delay = backoff; backoff = next_backoff(backoff); delay } } };
tokio::select! { _ = self.wake.notified() => {} _ = tokio::time::sleep(delay) => {} } } }
/// Performs one connection attempt, unless someone else already did. async fn connect(&self) -> ConnectionResult<()> { let _guard = self.attempt.lock().await;
if self.manager.is_connected().await { return Ok(()); }
self.manager.connect(&self.cookie).await?;
let reconnected = self.connects.fetch_add(1, Ordering::SeqCst) > 0; tracing::info!("Connected to {}", self.node());
if let Err(e) = logs::install(&self.manager).await { tracing::warn!("Cannot capture logs from {}: {e}", self.node()); }
if reconnected { self.announce_reconnect().await; }
Ok(()) }
/// Logs [`RECONNECT_NOTICE`] on the node, so that the gap shows up in the /// buffer `get_logs` reads. async fn announce_reconnect(&self) { let notice = rpc::binary_from_str(RECONNECT_NOTICE);
if let Err(e) = rpc_call(&self.manager, "logger", "notice", vec![notice], None).await { tracing::debug!("Cannot announce the reconnect on {}: {e}", self.node()); } }}
#[cfg(test)]#[allow(clippy::unwrap_used, clippy::expect_used)]mod tests { use super::*;
#[test] fn a_backoff_doubles_up_to_a_ceiling() { assert_eq!(next_backoff(MIN_BACKOFF), MIN_BACKOFF * 2); assert_eq!(next_backoff(MAX_BACKOFF), MAX_BACKOFF); assert_eq!(next_backoff(MAX_BACKOFF / 2 + MIN_BACKOFF), MAX_BACKOFF); }
#[tokio::test] async fn a_node_that_is_not_there_fails_rather_than_hangs() { let manager = Arc::new(ConnectionManager::for_test("nothing_is_here@localhost")); let supervisor = Supervisor::new("cookie".to_string(), manager);
assert!(supervisor.ensure_connected().await.is_err()); assert_eq!(supervisor.connects(), 0); }}