fix remove worker after instantly agent is reconnected

crates/relayer/src/agent_listener.rs

@@ -16,6 +16,7 @@ pub trait AgentConnection<S: AgentSubConnection<R, W>, R: AsyncRead + Unpin, W:
     Send + Sync
 {
     fn domain(&self) -> String;
+    fn conn_id(&self) -> u64;
     async fn create_sub_connection(&mut self) -> Result<S, Box<dyn Error>>;
     async fn recv(&mut self) -> Result<S, Box<dyn Error>>;
 }

crates/relayer/src/agent_listener/quic.rs

@@ -14,7 +14,7 @@
 use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
 use serde::de::DeserializeOwned;
 
-use crate::METRICS_AGENT_HISTOGRAM;
+use crate::{utils::now_ms, METRICS_AGENT_HISTOGRAM};
 
 use super::{AgentConnection, AgentListener, AgentSubConnection};
 
@@ -96,7 +96,11 @@
                     log::info!("register request domain {}", domain);
                     let res_buf = cluster_validator.sign_response_res(&request, None);
                     send.write_all(&res_buf).await?;
-                    Ok(AgentQuicConnection { domain, conn })
+                    Ok(AgentQuicConnection {
+                        domain,
+                        conn,
+                        conn_id: now_ms(),
+                    })
                 }
                 Err(e) => {
                     log::error!("invalid register request {:?}, error {}", request, e);
@@ -125,6 +129,7 @@
 
 pub struct AgentQuicConnection {
     domain: String,
+    conn_id: u64,
     conn: quinn::Connection,
 }
 
@@ -134,6 +139,10 @@
         self.domain.clone()
     }
 
+    fn conn_id(&self) -> u64 {
+        self.conn_id
+    }
+
     async fn create_sub_connection(&mut self) -> Result<AgentQuicSubConnection, Box<dyn Error>> {
         let (send, recv) = self.conn.open_bi().await?;
         Ok(AgentQuicSubConnection { send, recv })

crates/relayer/src/agent_listener/tcp.rs

@@ -16,7 +16,7 @@
 use protocol::key::ClusterValidator;
 use serde::de::DeserializeOwned;
 
-use crate::METRICS_AGENT_HISTOGRAM;
+use crate::{utils::now_ms, METRICS_AGENT_HISTOGRAM};
 
 use super::{AgentConnection, AgentListener, AgentSubConnection};
 
@@ -53,6 +53,7 @@
                     stream.write_all(&res).await?;
                     Ok(AgentTcpConnection {
                         domain,
+                        conn_id: now_ms(),
                         connection: yamux::Connection::new(
                             stream,
                             Default::default(),
@@ -111,6 +112,7 @@
 
 pub struct AgentTcpConnection {
     domain: String,
+    conn_id: u64,
     connection: yamux::Connection<TcpStream>,
 }
 
@@ -122,6 +124,10 @@
         self.domain.clone()
     }
 
+    fn conn_id(&self) -> u64 {
+        self.conn_id
+    }
+
     async fn create_sub_connection(&mut self) -> Result<AgentTcpSubConnection, Box<dyn Error>> {
         let client = OpenStreamsClient {
             connection: &mut self.connection,

crates/relayer/src/agent_store.rs

@@ -7,32 +7,44 @@
 
 use crate::ProxyTunnel;
 
+pub struct AgentEntry {
+    pub session: u64,
+    pub tx: Sender<Box<dyn ProxyTunnel>>,
+}
+
 #[derive(Clone, Default)]
 pub struct AgentStore {
     #[allow(clippy::type_complexity)]
-    agents: Arc<RwLock<HashMap<u64, Sender<Box<dyn ProxyTunnel>>>>>,
+    agents: Arc<RwLock<HashMap<u64, AgentEntry>>>,
 }
 
 impl AgentStore {
-    pub fn add(&self, id: u64, tx: Sender<Box<dyn ProxyTunnel>>) {
+    pub fn add(&self, id: u64, session: u64, tx: Sender<Box<dyn ProxyTunnel>>) {
         self.agents
             .write()
             .expect("Should write agents")
-            .insert(id, tx);
+            .insert(id, AgentEntry { tx, session });
     }
 
     pub fn get(&self, id: u64) -> Option<Sender<Box<dyn ProxyTunnel>>> {
         self.agents
             .read()
             .expect("Should write agents")
             .get(&id)
-            .cloned()
+            .map(|entry| entry.tx.clone())
     }
 
-    pub fn remove(&self, id: u64) {
-        self.agents
-            .write()
-            .expect("Should write agents")
-            .remove(&id);
+    pub fn remove(&self, id: u64, session: u64) -> bool {
+        let mut storage = self.agents.write().expect("Should write agents");
+
+        let current = storage.get(&id);
+        if let Some(entry) = current {
+            if entry.session == session {
+                storage.remove(&id);
+                return true;
+            }
+        }
+
+        false
     }
 }

crates/relayer/src/lib.rs

@@ -79,10 +79,12 @@
     counter!(METRICS_AGENT_COUNT).increment(1);
     log::info!("agent_connection.domain(): {}", agent_connection.domain());
     let domain = agent_connection.domain().to_string();
+    let conn_id = agent_connection.conn_id();
     let (mut agent_worker, proxy_tunnel_tx) =
         agent_worker::AgentWorker::<AG, S, R, W>::new(agent_connection, agent_rpc_handler);
     let home_id = home_id_from_domain(&domain);
-    agents.add(home_id, proxy_tunnel_tx);
+    log::info!("added home_id: {}, conn_id: {}", home_id, conn_id);
+    agents.add(home_id, conn_id, proxy_tunnel_tx);
     node_alias_sdk.register_alias(home_id).await;
     let agents = agents.clone();
     gauge!(METRICS_AGENT_LIVE).increment(1.0);
@@ -96,10 +98,16 @@
             }
         }
     }
-    agents.remove(home_id);
-    node_alias_sdk
-        .unregister_alias(home_id_from_domain(&domain))
-        .await;
+    let removed = agents.remove(home_id, conn_id);
+    log::info!(
+        "remove home_id: {}, conn_id: {}, removed: {}",
+        home_id,
+        conn_id,
+        removed
+    );
+    if removed {
+        node_alias_sdk.unregister_alias(home_id).await;
+    }
     log::info!("agent_worker exit for domain: {}", domain);
     gauge!(METRICS_AGENT_LIVE).decrement(1.0);
 }

crates/relayer/src/utils.rs

@@ -1,6 +1,7 @@
 use std::{
     collections::hash_map::DefaultHasher,
     hash::{Hash, Hasher},
+    time::{SystemTime, UNIX_EPOCH},
 };
 
 /// get home id from domain by get first subdomain
@@ -10,3 +11,11 @@
     parts.next().unwrap_or(domain).hash(&mut hasher);
     hasher.finish()
 }
+
+pub fn now_ms() -> u64 {
+    let start = SystemTime::now();
+    start
+        .duration_since(UNIX_EPOCH)
+        .expect("Time went backwards")
+        .as_millis() as u64
+}