Mirrord connection reconnect

changelog.d/+2901.fixed.md

@@ -0,0 +1 @@
+Add reconnection mechanism when using mirrord operator.

changelog.d/+update-single-portforward.changed.md

@@ -0,0 +1 @@
+Add a naive update to our port forward wrapper to force first check error channel instead of ranomly picking branch on `tokio::select!` impl.

mirrord/intproxy/Cargo.toml

@@ -37,6 +37,7 @@ hyper-util.workspace = true
 http-body-util.workspace = true
 bytes.workspace = true
 rand.workspace = true
+tokio-retry = "0.3"
 tokio-rustls.workspace = true
 rustls.workspace = true
 rustls-pemfile.workspace = true

mirrord/intproxy/src/agent_conn.rs

@@ -6,21 +6,20 @@ use std::{
     io,
     io::BufReader,
     net::{IpAddr, SocketAddr},
+    ops::ControlFlow,
     path::{Path, PathBuf},
-    sync::Arc,
+    sync::{Arc, LazyLock},
 };
 
-use mirrord_analytics::Reporter;
+use mirrord_analytics::{NullReporter, Reporter};
 use mirrord_config::LayerConfig;
 use mirrord_kube::{
-    api::{
-        kubernetes::{AgentKubernetesConnectInfo, KubernetesAPI},
-        wrap_raw_connection,
-    },
+    api::{kubernetes::AgentKubernetesConnectInfo, wrap_raw_connection},
     error::KubeApiError,
 };
 use mirrord_operator::client::{error::OperatorApiError, OperatorApi, OperatorSession};
 use mirrord_protocol::{ClientMessage, DaemonMessage};
+use semver::VersionReq;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::{
@@ -30,14 +29,26 @@ use tokio::{
         mpsc::{Receiver, Sender},
     },
 };
+use tokio_retry::{
+    strategy::{jitter, ExponentialBackoff},
+    Retry,
+};
 use tokio_rustls::TlsConnector;
 use tracing::Level;
 
 use crate::{
-    background_tasks::{BackgroundTask, MessageBus},
-    ProxyMessage,
+    background_tasks::{BackgroundTask, MessageBus, RestartableBackgroundTask},
+    main_tasks::{ConnectionRefresh, ProxyMessage},
 };
 
+mod portforward;
+
+static OPERATOR_RETRY_VERSION: LazyLock<VersionReq> = LazyLock::new(|| {
+    ">3.104.2"
+        .parse()
+        .expect("OPERATOR_RETRY_VERSION should be a valid VersionReq value")
+});
+
 /// Errors that can occur when the internal proxy tries to establish a connection with the agent.
 #[derive(Error, Debug)]
 pub enum AgentConnectionError {
@@ -97,6 +108,17 @@ pub enum AgentConnectInfo {
     DirectKubernetes(AgentKubernetesConnectInfo),
 }
 
+#[derive(Debug, Default)]
+pub enum ReconnectFlow {
+    ConnectInfo {
+        config: LayerConfig,
+        connect_info: AgentConnectInfo,
+    },
+
+    #[default]
+    Break,
+}
+
 /// Handles logic of the `proxy <-> agent` connection as a [`BackgroundTask`].
 ///
 /// # Note
@@ -107,6 +129,7 @@ pub enum AgentConnectInfo {
 pub struct AgentConnection {
     pub agent_tx: Sender<ClientMessage>,
     pub agent_rx: Receiver<DaemonMessage>,
+    pub reconnect: ReconnectFlow,
 }
 
 impl AgentConnection {
@@ -117,11 +140,23 @@ impl AgentConnection {
         connect_info: Option<AgentConnectInfo>,
         analytics: &mut R,
     ) -> Result<Self, AgentConnectionError> {
-        let (agent_tx, agent_rx) = match connect_info {
+        let (agent_tx, agent_rx, reconnect) = match connect_info {
             Some(AgentConnectInfo::Operator(session)) => {
                 let connection =
-                    OperatorApi::connect_in_existing_session(config, session, analytics).await?;
-                (connection.tx, connection.rx)
+                    OperatorApi::connect_in_existing_session(config, session.clone(), analytics)
+                        .await?;
+                (
+                    connection.tx,
+                    connection.rx,
+                    if OPERATOR_RETRY_VERSION.matches(&session.operator_version) {
+                        ReconnectFlow::ConnectInfo {
+                            config: config.clone(),
+                            connect_info: AgentConnectInfo::Operator(session),
+                        }
+                    } else {
+                        ReconnectFlow::default()
+                    },
+                )
             }
 
             Some(AgentConnectInfo::ExternalProxy(proxy_addr)) => {
@@ -131,7 +166,7 @@ impl AgentConnection {
 
                 let stream = socket.connect(proxy_addr).await?;
 
-                if config.external_proxy.tls_enable
+                let (tx, rx) = if config.external_proxy.tls_enable
                     && let (
                         Some(tls_certificate),
                         Some(client_tls_certificate),
@@ -140,8 +175,7 @@ impl AgentConnection {
                         config.external_proxy.tls_certificate.as_ref(),
                         config.internal_proxy.client_tls_certificate.as_ref(),
                         config.internal_proxy.client_tls_key.as_ref(),
-                    )
-                {
+                    ) {
                     wrap_connection_with_tls(
                         stream,
                         proxy_addr.ip(),
@@ -152,20 +186,14 @@ impl AgentConnection {
                     .await?
                 } else {
                     wrap_raw_connection(stream)
-                }
+                };
+
+                (tx, rx, ReconnectFlow::default())
             }
 
             Some(AgentConnectInfo::DirectKubernetes(connect_info)) => {
-                let k8s_api = KubernetesAPI::create(config)
-                    .await
-                    .map_err(AgentConnectionError::Kube)?;
-
-                let stream = k8s_api
-                    .create_connection_portforward(connect_info.clone())
-                    .await
-                    .map_err(AgentConnectionError::Kube)?;
-
-                wrap_raw_connection(stream)
+                let (tx, rx) = portforward::create_connection(config, connect_info.clone()).await?;
+                (tx, rx, ReconnectFlow::default())
             }
 
             None => {
@@ -174,18 +202,27 @@ impl AgentConnection {
                     .as_ref()
                     .ok_or(AgentConnectionError::NoConnectionMethod)?;
                 let stream = TcpStream::connect(address).await?;
-                wrap_raw_connection(stream)
+                let (tx, rx) = wrap_raw_connection(stream);
+                (tx, rx, ReconnectFlow::default())
             }
         };
 
-        Ok(Self { agent_tx, agent_rx })
+        Ok(Self {
+            agent_tx,
+            agent_rx,
+            reconnect,
+        })
     }
 
     pub async fn new_for_raw_address(address: SocketAddr) -> Result<Self, AgentConnectionError> {
         let stream = TcpStream::connect(address).await?;
         let (agent_tx, agent_rx) = wrap_raw_connection(stream);
 
-        Ok(Self { agent_tx, agent_rx })
+        Ok(Self {
+            agent_tx,
+            agent_rx,
+            reconnect: ReconnectFlow::Break,
+        })
     }
 
     #[tracing::instrument(level = Level::TRACE, name = "send_agent_message", skip(self), ret)]
@@ -205,7 +242,7 @@ impl BackgroundTask for AgentConnection {
     type MessageIn = ClientMessage;
     type MessageOut = ProxyMessage;
 
-    async fn run(mut self, message_bus: &mut MessageBus<Self>) -> Result<(), Self::Error> {
+    async fn run(&mut self, message_bus: &mut MessageBus<Self>) -> Result<(), Self::Error> {
         loop {
             tokio::select! {
                 msg = message_bus.recv() => match msg {
@@ -227,6 +264,58 @@ impl BackgroundTask for AgentConnection {
                         break Err(AgentChannelError);
                     }
                     Some(msg) => message_bus.send(ProxyMessage::FromAgent(msg)).await,
+                },
+            }
+        }
+    }
+}
+
+impl RestartableBackgroundTask for AgentConnection {
+    #[tracing::instrument(level = Level::TRACE, skip(self, message_bus), ret)]
+    async fn restart(
+        &mut self,
+        error: Self::Error,
+        message_bus: &mut MessageBus<Self>,
+    ) -> ControlFlow<Self::Error> {
+        match &self.reconnect {
+            ReconnectFlow::Break => ControlFlow::Break(error),
+            ReconnectFlow::ConnectInfo {
+                config,
+                connect_info,
+            } => {
+                message_bus
+                    .send(ProxyMessage::ConnectionRefresh(ConnectionRefresh::Start))
+                    .await;
+
+                let retry_strategy = ExponentialBackoff::from_millis(50).map(jitter).take(10);
+
+                let connection = Retry::spawn(retry_strategy, || async move {
+                    AgentConnection::new(
+                        config,
+                        connect_info.clone().into(),
+                        &mut NullReporter::default(),
+                    )
+                    .await
+                    .inspect_err(
+                        |err| tracing::error!(error = ?err, "unable to connect to agent upon retry"),
+                    )
+                })
+                .await;
+
+                match connection {
+                    Ok(connection) => {
+                        *self = connection;
+                        message_bus
+                            .send(ProxyMessage::ConnectionRefresh(ConnectionRefresh::End))
+                            .await;
+
+                        ControlFlow::Continue(())
+                    }
+                    Err(error) => {
+                        tracing::error!(?error, "unable to reconnect agent");
+
+                        ControlFlow::Break(AgentChannelError)
+                    }
                 }
             }
         }

mirrord/intproxy/src/agent_conn/portforward.rs

@@ -0,0 +1,25 @@
+use mirrord_config::LayerConfig;
+use mirrord_kube::api::{
+    kubernetes::{AgentKubernetesConnectInfo, KubernetesAPI},
+    wrap_raw_connection,
+};
+use mirrord_protocol::{ClientMessage, DaemonMessage};
+use tokio::sync::mpsc;
+
+use crate::agent_conn::AgentConnectionError;
+
+pub async fn create_connection(
+    config: &LayerConfig,
+    connect_info: AgentKubernetesConnectInfo,
+) -> Result<(mpsc::Sender<ClientMessage>, mpsc::Receiver<DaemonMessage>), AgentConnectionError> {
+    let k8s_api = KubernetesAPI::create(config)
+        .await
+        .map_err(AgentConnectionError::Kube)?;
+
+    let stream = k8s_api
+        .create_connection_portforward(connect_info.clone())
+        .await
+        .map_err(AgentConnectionError::Kube)?;
+
+    Ok(wrap_raw_connection(stream))
+}