-
Notifications
You must be signed in to change notification settings - Fork 1
fix(stream): enable public reconnect() with needs_reconnect() flag #2
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -57,6 +57,8 @@ pub struct WebSocketStream { | |
| needs_pong_flush: bool, | ||
| /// Pending messages from array-formatted snapshot (e.g. initial book snapshot) | ||
| pending_books: VecDeque<StreamMessage>, | ||
| /// Flag indicating reconnection is needed (set when connection lost) | ||
| needs_reconnect: bool, | ||
| } | ||
|
|
||
| /// Stream statistics | ||
|
|
@@ -113,6 +115,7 @@ impl WebSocketStream { | |
| reconnect_config: ReconnectConfig::default(), | ||
| needs_pong_flush: false, | ||
| pending_books: VecDeque::new(), | ||
| needs_reconnect: false, | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -492,9 +495,16 @@ impl WebSocketStream { | |
| } | ||
| } | ||
|
|
||
| /// Check if reconnection is needed (connection was lost) | ||
| pub fn needs_reconnect(&self) -> bool { | ||
| self.needs_reconnect | ||
| } | ||
|
|
||
| /// Reconnect with exponential backoff | ||
| #[allow(dead_code)] | ||
| async fn reconnect(&mut self) -> Result<()> { | ||
| /// | ||
| /// Call this method when `needs_reconnect()` returns true. | ||
| /// It will attempt to reconnect to the WebSocket and resubscribe to all previous subscriptions. | ||
| pub async fn reconnect(&mut self) -> Result<()> { | ||
|
Comment on lines
495
to
+507
|
||
| let mut delay = self.reconnect_config.base_delay; | ||
| let mut retries = 0; | ||
|
|
||
|
|
@@ -505,6 +515,7 @@ impl WebSocketStream { | |
| Ok(()) => { | ||
| info!("Successfully reconnected"); | ||
| self.stats.reconnect_count += 1; | ||
| self.needs_reconnect = false; // Reset flag on success | ||
|
|
||
| // Resubscribe to all previous subscriptions | ||
| let subscriptions = self.subscriptions.clone(); | ||
|
|
@@ -688,6 +699,7 @@ impl Stream for WebSocketStream { | |
| tokio_tungstenite::tungstenite::Message::Close(_) => { | ||
| info!("WebSocket connection closed by server"); | ||
| self.connection = None; | ||
| self.needs_reconnect = true; | ||
| Poll::Ready(None) | ||
| }, | ||
| _ => Poll::Pending, | ||
|
|
@@ -696,16 +708,28 @@ impl Stream for WebSocketStream { | |
| Poll::Ready(Some(Err(e))) => { | ||
| error!("WebSocket error: {}", e); | ||
| self.stats.errors += 1; | ||
| self.needs_reconnect = true; | ||
| self.connection = None; | ||
| Poll::Ready(Some(Err(e.into()))) | ||
| }, | ||
| Poll::Ready(None) => { | ||
| debug!("WebSocket stream ended"); | ||
| self.needs_reconnect = true; | ||
| self.connection = None; | ||
| Poll::Ready(None) | ||
| }, | ||
| Poll::Pending => Poll::Pending, | ||
| } | ||
| } else { | ||
| Poll::Ready(None) | ||
| // Connection is None - check if we're waiting for reconnection | ||
| if self.needs_reconnect { | ||
| // Return Pending to avoid busy loop in select! while waiting for reconnect() | ||
| // The consumer should call reconnect() when they see needs_reconnect() | ||
| Poll::Pending | ||
| } else { | ||
| // Stream is truly ended (no reconnect requested) | ||
| Poll::Ready(None) | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
needs_reconnectis only reset insidereconnect(). If the stream re-establishes a connection via another path that callsconnect()(e.g.,subscribe_asyncwhenconnection.is_none()),needs_reconnect()can continue to returntrueeven after a successful connect. Consider resetting the flag on any successful (re)connection (e.g., inconnect()), or ensuring all reconnection entry points funnel throughreconnect().