Implemented heartbeat

Author: Linwenxuan04

crates/lagrange-core-runner/src/main.rs

@@ -9,7 +9,7 @@ const BOT_PROTOCOL: Protocols = Protocols::Linux;
 const BOT_VERBOSE: bool = false;
 const AUTO_RECONNECT: bool = true;
 const AUTO_RELOGIN: bool = true;
-const LOG_LEVEL: Level = Level::INFO;
+const LOG_LEVEL: Level = Level::TRACE;
 const COLORED_LOGS: bool = true;
 
 #[tokio::main]
@@ -49,6 +49,11 @@ async fn main() -> Result<()> {
     info!("Starting main event loop...");
     info!("Press Ctrl+C to shutdown gracefully");
 
+    context.connect().await.expect("Failed to establish initial connection");
+
+    // Start connection monitor for auto-reconnect
+    context.clone().start_connection_monitor();
+
     match tokio::signal::ctrl_c().await {
         Ok(()) => {
             info!("Received shutdown signal, cleaning up...");
@@ -65,7 +70,8 @@ async fn main() -> Result<()> {
 
 fn setup_tracing() -> Result<()> {
     let env_filter =
-        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(LOG_LEVEL.as_str()));
+        EnvFilter::try_from_default_env()
+            .unwrap_or_else(|_| EnvFilter::new(LOG_LEVEL.as_str()));
 
     let fmt_layer = fmt::layer()
         .with_target(true)

crates/lagrange-core/src/business.rs

@@ -0,0 +1 @@
+pub mod network;

crates/lagrange-core/src/business/network.rs

@@ -0,0 +1,97 @@
+use crate::{BotContext, Error, internal::services::AliveEventReq};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::time;
+
+impl BotContext {
+    pub async fn connect(self: &Arc<Self>) -> Result<bool, Error> {
+        let result = self.socket.connect(
+            self.config.use_ipv6_network,
+            self.packet.clone()
+        ).await;
+
+        if result.is_err() {
+            Err(Error::NetworkError("Failed to connect to server".to_string()))
+        } else {
+            self.clone().start_heartbeat();
+            Ok(true)
+        }
+    }
+
+    /// Start sending heartbeat packets at 5-second intervals
+    pub fn start_heartbeat(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
+        tokio::spawn(async move {
+            let mut interval = time::interval(Duration::from_secs(5));
+            interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
+
+            loop {
+                interval.tick().await;
+
+                if !self.socket.is_connected().await {
+                    tracing::debug!("Socket not connected, skipping heartbeat");
+                    continue;
+                }
+
+                if let Err(e) = self.event.send_event(self.clone(), AliveEventReq {}).await {
+                    tracing::warn!(error = %e, "Failed to send heartbeat");
+                }
+            }
+        })
+    }
+
+    pub fn start_connection_monitor(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
+        tokio::spawn(async move {
+            if !self.config.auto_reconnect {
+                tracing::info!("Auto-reconnect disabled, connection monitor not started");
+                return;
+            }
+
+            tracing::info!("Starting connection monitor with auto-reconnect enabled");
+            let mut check_interval = time::interval(Duration::from_secs(3));
+            check_interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
+
+            let mut retry_count = 0u32;
+            let max_backoff_secs = 60; // Max 60 seconds between retries
+
+            loop {
+                check_interval.tick().await;
+
+                if self.socket.is_connected().await {
+                    if retry_count > 0 {
+                        retry_count = 0;
+                        tracing::info!("Connection restored, retry count reset");
+                    }
+                    continue;
+                }
+
+                tracing::warn!(retry_count, "Socket disconnected, attempting to reconnect");
+                let backoff_secs = (1u64 << retry_count.min(6)).min(max_backoff_secs);
+
+                if retry_count > 0 {
+                    tracing::info!(backoff_secs, "Waiting before reconnection attempt");
+                    time::sleep(Duration::from_secs(backoff_secs)).await;
+                }
+
+                match self.socket.connect(
+                    self.config.use_ipv6_network,
+                    self.packet.clone()
+                ).await {
+                    Ok(_) => {
+                        tracing::info!("Successfully reconnected to server");
+                        self.clone().start_heartbeat();
+                        retry_count = 0;
+                    }
+                    Err(e) => {
+                        retry_count += 1;
+                        tracing::error!(
+                            error = %e,
+                            retry_count,
+                            next_backoff_secs = (1u64 << retry_count.min(6)).min(max_backoff_secs),
+                            "Failed to reconnect"
+                        );
+                    }
+                }
+            }
+        })
+    }
+}

crates/lagrange-core/src/context.rs

@@ -116,7 +116,6 @@ impl BotContextBuilder {
 
         let cache = CacheContext::new();
         let socket = SocketContext::new();
-        let event = EventContext::new();
 
         let keystore_arc = Arc::new(std::sync::RwLock::new(keystore.clone()));
         let app_info_arc = Arc::new(app_info.clone());
@@ -126,6 +125,10 @@ impl BotContextBuilder {
 
         let service = ServiceContext::new(&config);
 
+        // EventContext needs service, packet, socket, and config
+        let config_arc = Arc::new(config.clone());
+        let event = EventContext::new(service.clone(), packet.clone(), socket.clone(), config_arc);
+
         Arc::new(BotContext {
             config,
             app_info,

crates/lagrange-core/src/internal.rs

@@ -1,3 +1,6 @@
 pub mod context;
 mod packets;
 pub mod services;
+
+// Re-export commonly used packet types
+pub use packets::SsoPacket;

crates/lagrange-core/src/internal/context/event.rs

@@ -1,15 +1,32 @@
 use crate::protocol::{EventMessage, ProtocolEvent};
+use crate::config::BotConfig;
+use super::{PacketContext, ServiceContext, SocketContext};
 use std::{any::TypeId, sync::Arc};
 use tokio::sync::broadcast;
 
 pub struct EventContext {
     sender: broadcast::Sender<EventMessage>,
+    service: Arc<ServiceContext>,
+    packet: Arc<PacketContext>,
+    socket: Arc<SocketContext>,
+    config: Arc<BotConfig>,
 }
 
 impl EventContext {
-    pub fn new() -> Arc<Self> {
+    pub fn new(
+        service: Arc<ServiceContext>,
+        packet: Arc<PacketContext>,
+        socket: Arc<SocketContext>,
+        config: Arc<BotConfig>,
+    ) -> Arc<Self> {
         let (sender, _) = broadcast::channel(1024);
-        Arc::new(Self { sender })
+        Arc::new(Self {
+            sender,
+            service,
+            packet,
+            socket,
+            config,
+        })
     }
 
     pub fn post_event(&self, event: EventMessage) {
@@ -31,6 +48,49 @@ impl EventContext {
             _phantom: std::marker::PhantomData,
         }
     }
+
+    /// Send a protocol event as a packet through the network and wait for response
+    pub async fn send_event<T>(
+        self: &Arc<Self>,
+        context: Arc<crate::context::BotContext>,
+        event: T,
+    ) -> Result<(), crate::Error>
+    where
+        T: ProtocolEvent,
+    {
+        use crate::internal::context::packet::ServiceAttribute;
+
+        let event_msg = EventMessage::new(event);
+        let bytes = self.service.resolve_outgoing(event_msg.clone(), context.clone()).await?;
+
+        let event_type = event_msg.type_id();
+        let mappings = crate::internal::services::registry()
+            .get_event_mappings(event_type)
+            .ok_or_else(|| crate::Error::ServiceNotFound(format!("event type {:?}", event_type)))?;
+
+        let service = mappings
+            .iter()
+            .find(|m| (self.config.protocol as u8) & m.protocol != 0)
+            .ok_or_else(|| crate::Error::ServiceNotFound(format!("No service for protocol {:?}", self.config.protocol)))?;
+
+        let metadata = service.service.metadata();
+
+        let attributes = Some(ServiceAttribute::new()
+            .with_request_type(metadata.request_type)
+            .with_encrypt_type(metadata.encrypt_type));
+
+        let response = self.packet.send_packet(
+            metadata.command.to_string(),
+            bytes,
+            self.socket.clone(),
+            attributes,
+        ).await?;
+
+        let response_event = self.service.resolve_incoming(&response, context).await?;
+        self.post_event(response_event);
+
+        Ok(())
+    }
 }
 
 pub struct TypedEventReceiver<T> {
@@ -63,12 +123,6 @@ impl<T: 'static> TypedEventReceiver<T> {
     }
 }
 
-impl Default for EventContext {
-    fn default() -> Self {
-        let (sender, _) = broadcast::channel(1024);
-        Self { sender }
-    }
-}
 
 impl<T> Clone for TypedEventReceiver<T> {
     fn clone(&self) -> Self {

crates/lagrange-core/src/internal/context/packet.rs

@@ -3,10 +3,10 @@ use crate::{
     config::BotConfig,
     error::{Error, Result},
     internal::packets::{
-        EncryptType, RequestType, ServicePacker, SsoPacker, SsoPacket, SsoSecureInfo,
+        ServicePacker, SsoPacker, SsoPacket, SsoSecureInfo,
     },
     keystore::BotKeystore,
-    protocol::Protocols,
+    protocol::{EncryptType, Protocols, RequestType},
 };
 use bytes::Bytes;
 use dashmap::DashMap;
@@ -95,11 +95,23 @@ impl PacketContext {
             extra: String::new(),
         };
 
+        tracing::debug!(
+            sequence_u32 = sequence,
+            sequence_i32 = sso_packet.sequence,
+            command = %command,
+            "Sending packet and registering pending task"
+        );
+
         let encoded = self.encode_packet(&sso_packet, attributes).await?;
 
-        socket.send(encoded)?;
+        socket.send(encoded).await?;
 
         let response = rx.await.map_err(|_| {
+            tracing::warn!(
+                sequence = sequence,
+                command = %command,
+                "Response channel closed, removing pending task"
+            );
             self.pending_tasks.remove(&sequence);
             Error::NetworkError("Response channel closed".to_string())
         })?;
@@ -110,6 +122,13 @@ impl PacketContext {
     pub fn dispatch_packet(&self, packet: SsoPacket) -> Option<SsoPacket> {
         let sequence = packet.sequence as u32;
 
+        tracing::debug!(
+            packet_sequence_i32 = packet.sequence,
+            converted_sequence_u32 = sequence,
+            pending_tasks_count = self.pending_tasks.len(),
+            "Attempting to dispatch packet"
+        );
+
         if let Some((_, sender)) = self.pending_tasks.remove(&sequence) {
             if packet.ret_code != 0 {
                 tracing::error!(
@@ -121,14 +140,33 @@ impl PacketContext {
                 );
             }
 
+            tracing::debug!(
+                sequence = sequence,
+                command = %packet.command,
+                "Successfully matched and removed pending task"
+            );
+
             let _ = sender.send(packet);
             None
         } else {
+            // Collect all pending sequence numbers for debugging
+            let pending_sequences: Vec<u32> = self.pending_tasks.iter()
+                .map(|entry| *entry.key())
+                .collect();
+
+            tracing::warn!(
+                sequence_i32 = packet.sequence,
+                sequence_u32 = sequence,
+                command = %packet.command,
+                pending_tasks_count = self.pending_tasks.len(),
+                pending_sequences = ?pending_sequences,
+                "Failed to find pending task for sequence - packet will be routed to services"
+            );
             Some(packet)
         }
     }
 
-    async fn encode_packet(
+    pub async fn encode_packet(
         &self,
         packet: &SsoPacket,
         attributes: Option<ServiceAttribute>,