mark packet_router's shutdown as success before drop (#4491)

nymtech · Mar 25, 2024 · e352c25 · e352c25
1 parent f5378e8
commit e352c25
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/common/client-libs/gateway-client/src/packet_router.rs b/common/client-libs/gateway-client/src/packet_router.rs
@@ -69,6 +69,10 @@ impl PacketRouter {
         }
         Ok(())
     }
+
+    pub fn mark_as_success(&mut self) {
+        self.shutdown.mark_as_success();
+    }
 }
 
 impl GatewayPacketRouter for PacketRouter {

diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs
@@ -97,7 +97,7 @@ impl PartiallyDelegated {
 
     pub(crate) fn split_and_listen_for_mixnet_messages(
         conn: WsConn,
-        packet_router: PacketRouter,
+        mut packet_router: PacketRouter,
         shared_key: Arc<SharedKeys>,
         mut shutdown: TaskClient,
     ) -> Self {
@@ -140,6 +140,7 @@ impl PartiallyDelegated {
             if match ret_err {
                 Err(err) => stream_sender.send(Err(err)),
                 Ok(_) => {
+                    packet_router.mark_as_success();
                     shutdown.mark_as_success();
                     stream_sender.send(Ok(stream))
                 }