Inner select on on_messages

* Send shutdown instead of panic when reaching max fail

* Stop quicker on failure

* Update comment

common/client-core/src/client/mix_traffic/mod.rs

@@ -103,15 +103,23 @@ impl MixTrafficController {
                 tokio::select! {
                     mix_packets = self.mix_rx.recv() => match mix_packets {
                         Some(mix_packets) => {
-                            if let Err(err) = self.on_messages(mix_packets).await {
-                                error!("Failed to send sphinx packet(s) to the gateway: {err}");
-                                if self.consecutive_gateway_failure_count == MAX_FAILURE_COUNT {
-                                    // Disconnect from the gateway. If we should try to re-connect
-                                    // is handled at a higher layer.
-                                    error!("Failed to send sphinx packet to the gateway {MAX_FAILURE_COUNT} times in a row - assuming the gateway is dead");
-                                    // Do we need to handle the embedded mixnet client case
-                                    // separately?
-                                    shutdown.send_we_stopped(Box::new(ClientCoreError::GatewayFailedToForwardMessages));
+                            tokio::select! {
+                                ret = self.on_messages(mix_packets) => {
+                                    if let Err(err) = ret {
+                                        error!("Failed to send sphinx packet(s) to the gateway: {err}");
+                                        if self.consecutive_gateway_failure_count == MAX_FAILURE_COUNT {
+                                            // Disconnect from the gateway. If we should try to re-connect
+                                            // is handled at a higher layer.
+                                            error!("Failed to send sphinx packet to the gateway {MAX_FAILURE_COUNT} times in a row - assuming the gateway is dead");
+                                            // Do we need to handle the embedded mixnet client case
+                                            // separately?
+                                            shutdown.send_we_stopped(Box::new(ClientCoreError::GatewayFailedToForwardMessages));
+                                            break;
+                                        }
+                                    }
+                                }
+                                _ = shutdown.recv_with_delay() => {
+                                    log::trace!("MixTrafficController: Received shutdown");
                                     break;
                                 }
                             }

common/client-libs/gateway-client/src/error.rs

@@ -118,6 +118,9 @@ pub enum GatewayClientError {
         "this operation couldn't be completed as the program is in the process of shutting down"
     )]
     ShutdownInProgress,
+
+    #[error("Timed out on sending message(s)")]
+    TimeoutOnSink,
 }
 
 impl GatewayClientError {

common/client-libs/gateway-client/src/socket_state.rs

@@ -291,7 +291,10 @@ impl PartiallyDelegatedHandle {
         &mut self,
         msg: Message,
     ) -> Result<(), GatewayClientError> {
-        Ok(self.sink_half.send(msg).await?)
+        tokio::time::timeout(std::time::Duration::from_secs(1), self.sink_half.send(msg))
+            .await
+            .map_err(|_| GatewayClientError::TimeoutOnSink)??;
+        Ok(())
     }
 
     pub(crate) async fn batch_send_without_response(
@@ -300,7 +303,13 @@ impl PartiallyDelegatedHandle {
     ) -> Result<(), GatewayClientError> {
         let stream_messages: Vec<_> = messages.into_iter().map(Ok).collect();
         let mut send_stream = futures::stream::iter(stream_messages);
-        Ok(self.sink_half.send_all(&mut send_stream).await?)
+        tokio::time::timeout(
+            std::time::Duration::from_secs(1),
+            self.sink_half.send_all(&mut send_stream),
+        )
+        .await
+        .map_err(|_| GatewayClientError::TimeoutOnSink)??;
+        Ok(())
     }
 
     pub(crate) async fn merge(self) -> Result<WsConn, GatewayClientError> {