fix(gateway): fix gc logic

Author: MasterPtato

engine/packages/pegboard-gateway/src/shared_state.rs

@@ -31,6 +31,7 @@ struct InFlightRequest {
 	opened: bool,
 	pending_msgs: Vec<PendingMessage>,
 	hibernation_state: Option<HibernationState>,
+	stopping: bool,
 }
 
 pub struct PendingMessage {
@@ -98,13 +99,19 @@ impl SharedState {
 					opened: false,
 					pending_msgs: Vec::new(),
 					hibernation_state: None,
+					stopping: false,
 				});
 			}
 			Entry::Occupied(mut entry) => {
 				entry.receiver_subject = receiver_subject;
 				entry.msg_tx = msg_tx;
 				entry.opened = false;
 				entry.pending_msgs.clear();
+
+				if entry.stopping {
+					entry.hibernation_state = None;
+					entry.stopping = false;
+				}
 			}
 		}
 
@@ -403,6 +410,8 @@ impl SharedState {
 	async fn gc(&self) {
 		#[derive(Debug)]
 		enum MsgGcReason {
+			/// Gateway channel is closed and there are no pending messages
+			GatewayClosed,
 			/// Any tunnel message not acked (TunnelAck)
 			MessageNotAcked,
 			/// WebSocket pending messages (ToServerWebSocketMessageAck)
@@ -417,13 +426,30 @@ impl SharedState {
 
 			let now = Instant::now();
 
+			// First, check if an in flight req is beyond the timeout for tunnel message ack and websocket
+			// message ack
 			self.in_flight_requests
-				.retain_async(|request_id, req| {
-					if req.msg_tx.is_closed() {
-						return false;
+				.iter_mut_async(|mut entry| {
+					let (request_id, req) = &mut *entry;
+
+					if req.stopping {
+						return true;
 					}
 
 					let reason = 'reason: {
+						// If we have no pending messages of any kind and the channel is closed, remove the
+						// in flight req
+						if req.msg_tx.is_closed()
+							&& req.pending_msgs.is_empty()
+							&& req
+								.hibernation_state
+								.as_ref()
+								.map(|hs| hs.pending_ws_msgs.is_empty())
+								.unwrap_or(true)
+						{
+							break 'reason Some(MsgGcReason::GatewayClosed);
+						}
+
 						if let Some(earliest_pending_msg) = req.pending_msgs.first() {
 							if now.duration_since(earliest_pending_msg.send_instant)
 								<= MESSAGE_ACK_TIMEOUT
@@ -449,13 +475,34 @@ impl SharedState {
 						tracing::debug!(
 							request_id=?Uuid::from_bytes(*request_id),
 							?reason,
-							"gc collecting in flight request"
+							"gc stopping in flight request"
 						);
+
 						let _ = req.msg_tx.send(TunnelMessageData::Timeout);
+
+						// Mark req as stopping to skip this loop next time the gc is run
+						req.stopping = true;
+					}
+
+					true
+				})
+				.await;
+
+			self.in_flight_requests
+				.retain_async(|request_id, req| {
+					// The reason we check for stopping here is because msg_rx could be dropped if we are
+					// between websocket retries (we don't want to remove the in flight req in this case).
+					// When the websocket reconnects a new channel will be created
+					if req.stopping && req.msg_tx.is_closed() {
+						tracing::debug!(
+							request_id=?Uuid::from_bytes(*request_id),
+							"gc removing in flight request"
+						);
+
+						return false;
 					}
 
-					// Return true if the request was not gc'd
-					reason.is_none()
+					true
 				})
 				.await;
 		}