连接成功事件调用

This commit is contained in:
zqm
2026-04-07 16:09:34 +08:00
parent de73921fc2
commit fdbf19b722
3 changed files with 36 additions and 7 deletions

View File

@@ -30,4 +30,4 @@
pub mod websocket;
pub use websocket::{WebSocketClient, WebSocketConfig, WebSocketMessage, OutgoingMessage, ConnectionStatus};
pub use websocket::{WebSocketClient, WebSocketConfig, WebSocketMessage, OutgoingMessage, ConnectionStatus, ReconnectedCallback};

View File

@@ -58,7 +58,13 @@ pub type StatusCallback = Arc<dyn Fn(ConnectionStatus, ConnectionStatus) + Send
pub type SentCallback = Arc<dyn Fn(String, Value) + Send + Sync>;
/// Callback triggered before reconnecting
/// Arguments: (attempt_number, url_arc) - app can update the URL via url_arc
pub type ReconnectingCallback = Arc<dyn Fn(u32, Arc<tokio::sync::Mutex<String>>) + Send + Sync>;
/// Note: This is an async callback - return a boxed Future
pub type ReconnectingCallback = Arc<dyn Fn(u32, Arc<tokio::sync::Mutex<String>>) -> Pin<Box<dyn std::future::Future<Output = ()> + Send + Sync + 'static>> + Send + Sync>;
/// Callback triggered after successful reconnection (after the first connection)
/// Arguments: (url, send_fn) - send_fn can be called to send messages
/// Note: This is an async callback - return a boxed Future
pub type ReconnectedCallback = Arc<dyn Fn(String, Arc<Mutex<Option<mpsc::Sender<OutgoingMessage>>>>) -> Pin<Box<dyn std::future::Future<Output = ()> + Send + Sync + 'static>> + Send + Sync>;
/// Callback triggered on first successful connection (before any reconnect)
/// Arguments: (url, send_fn) - send_fn can be called to send messages
/// Note: This is an async callback - return a boxed Future
@@ -87,6 +93,8 @@ pub struct WebSocketClient {
on_reconnecting: Option<ReconnectingCallback>,
/// Callback triggered on first successful connection
on_first_connect: Option<FirstConnectCallback>,
/// Callback triggered after successful reconnection (after the first connection)
on_reconnected: Option<ReconnectedCallback>,
// Reconnection state
reconnect_attempts: Arc<Mutex<u32>>,
@@ -113,6 +121,7 @@ impl WebSocketClient {
on_message_sent: None,
on_reconnecting: None,
on_first_connect: None,
on_reconnected: None,
reconnect_attempts: Arc::new(Mutex::new(0)),
reconnect_delay_ms: Arc::new(Mutex::new(config.reconnect_delay_ms)),
}
@@ -189,11 +198,12 @@ impl WebSocketClient {
}
/// Set callback for reconnecting (called before each reconnect attempt)
/// This is an async callback - the returned Future will be awaited
/// The callback receives: (attempt_number, url_arc)
/// Use url_arc to update the URL for the next connection attempt
pub fn on_reconnecting<F>(&mut self, callback: F) -> &mut Self
where
F: Fn(u32, Arc<tokio::sync::Mutex<String>>) + Send + Sync + 'static,
F: Fn(u32, Arc<tokio::sync::Mutex<String>>) -> Pin<Box<dyn std::future::Future<Output = ()> + Send + Sync + 'static>> + Send + Sync + 'static,
{
self.on_reconnecting = Some(Arc::new(callback));
self
@@ -209,6 +219,16 @@ impl WebSocketClient {
self
}
/// Set callback for reconnection (called after successful reconnection, not on first connection)
/// This is an async callback - the returned Future will be awaited
pub fn on_reconnected<F>(&mut self, callback: F) -> &mut Self
where
F: Fn(String, Arc<Mutex<Option<mpsc::Sender<OutgoingMessage>>>>) -> Pin<Box<dyn std::future::Future<Output = ()> + Send + Sync + 'static>> + Send + Sync + 'static,
{
self.on_reconnected = Some(Arc::new(callback));
self
}
/// Update the URL for future reconnection attempts
/// Call this method when the server URL changes
pub fn update_url(&mut self, new_url: String) {
@@ -261,6 +281,7 @@ impl WebSocketClient {
let on_error = self.on_error.clone();
let on_reconnecting = self.on_reconnecting.clone();
let on_first_connect = self.on_first_connect.clone();
let on_reconnected = self.on_reconnected.clone();
// Spawn the WebSocket task
*self.is_running.lock().await = true;
@@ -283,6 +304,7 @@ impl WebSocketClient {
on_error,
on_reconnecting,
on_first_connect,
on_reconnected,
sender,
)
.await;
@@ -315,6 +337,7 @@ impl WebSocketClient {
on_error: Option<ErrorCallback>,
on_reconnecting: Option<ReconnectingCallback>,
on_first_connect: Option<FirstConnectCallback>,
on_reconnected: Option<ReconnectedCallback>,
sender: Arc<Mutex<Option<mpsc::Sender<OutgoingMessage>>>>,
) {
loop {
@@ -332,7 +355,7 @@ impl WebSocketClient {
// Clone sender for first_connect callback
let sender_clone = Arc::clone(&sender);
match Self::connect_and_handle(&current_url, &config, &mut receiver, &queue, &on_connected, &on_message, &on_binary, &on_error, is_first, &on_first_connect, sender_clone).await {
match Self::connect_and_handle(&current_url, &config, &mut receiver, &queue, &on_connected, &on_message, &on_binary, &on_error, is_first, &on_first_connect, &on_reconnected, sender_clone).await {
Ok(_) => {
if config.debug_mode {
debug!("WebSocket connection closed normally");
@@ -349,7 +372,7 @@ impl WebSocketClient {
// Trigger reconnecting callback
if let Some(ref callback) = on_reconnecting {
callback(attempts, client_url.clone());
callback(attempts, client_url.clone()).await;
}
*status.lock().await = ConnectionStatus::Reconnecting;
@@ -384,7 +407,7 @@ impl WebSocketClient {
// Trigger reconnecting callback (allows app to update URL, reload config, etc.)
if let Some(ref callback) = on_reconnecting {
callback(attempts, client_url.clone());
callback(attempts, client_url.clone()).await;
}
*status.lock().await = ConnectionStatus::Reconnecting;
@@ -420,6 +443,7 @@ impl WebSocketClient {
on_error: &Option<ErrorCallback>,
is_first: bool,
on_first_connect: &Option<FirstConnectCallback>,
on_reconnected: &Option<ReconnectedCallback>,
sender: Arc<Mutex<Option<mpsc::Sender<OutgoingMessage>>>>,
) -> Result<(), WebSocketError> {
use http::header::{HeaderName, HeaderValue};
@@ -472,6 +496,11 @@ impl WebSocketClient {
if let Some(ref callback) = on_first_connect {
callback(url.to_string(), sender).await;
}
} else {
// Call on_reconnected callback (after successful reconnection)
if let Some(ref callback) = on_reconnected {
callback(url.to_string(), sender).await;
}
}
// Start heartbeat task if enabled

View File

@@ -12,7 +12,7 @@ mod config;
mod message;
mod error;
pub use client::{WebSocketClient, ConnectionStatus, ReconnectingCallback, OutgoingMessage};
pub use client::{WebSocketClient, ConnectionStatus, ReconnectingCallback, ReconnectedCallback, OutgoingMessage};
pub use config::WebSocketConfig;
pub use message::{WebSocketMessage, ClientType};
pub use error::WebSocketError;