Axum get stuck when websocket session closed hardly

[doowonee]

Summary

Hello I'm making a websocket server with axum and I'm testing in real situation like below.

1. How many session can handled
2. How may message per second can handle.
3. What happen when some of session closed while broadcasting data to websocket sessions

It looks fine but there is only one problem. When below errors occurred the axum server never respond anything even html handler does not work. It just get stuck until restarting program.

# It is simulating number 3 situation
Io(Os { code: 54, kind: ConnectionReset, message: "Connection reset by peer" }) 
Io(Os { code: 32, kind: BrokenPipe, message: "Broken pipe" })

Can you give me some advice how to solve this? I use split function of futures just like an example of yours.

axum version

0.7.4

[doowonee]

I used https://github.com/tokio-rs/console to investigate about this and I found something.

As you see Some of tokio task get stuck on 3 points of my source code.

// just handle splitted stream read line 84
  let mut session_forward_task = tokio::spawn(async move {
      while let Some(m) = session_rx.recv().await {
          if let Err(e) = wb_tx.send(m).await {
              tracing::error!("session forward error: {e:?} {cloned_session_id}");
              return;
          }
      }
  });

// line 104
  let mut recv_task = tokio::spawn(async move {
      on_rx(cloned_epc, wb_rx, session).await;
  });

And the last one cargo>/axum-0.7.4/src/extract/ws.rs:313:9 maybe axum websocket extracter is the problem?

// https://github.com/tokio-rs/axum/blob/main/axum/src/extract/ws.rs#L313
// axum websocket extracter
  tokio::spawn(async move {
            let upgraded = match on_upgrade.await {
                Ok(upgraded) => upgraded,
                Err(err) => {
                    on_failed_upgrade.call(Error::new(err));
                    return;
                }
            };
            let upgraded = TokioIo::new(upgraded);

            let socket =
                WebSocketStream::from_raw_socket(upgraded, protocol::Role::Server, Some(config))
                    .await;
            let socket = WebSocket {
                inner: socket,
                protocol,
            };
            callback(socket).await;
        });