What causes "Failed to close WebSocket stream" in tracel-ai/burn?

Treat close failures as non-fatal: replace expect with a logged warning, since the socket is being discarded anyway.

Failed to close WebSocket stream

Error message

Failed to close WebSocket stream

What it means

This panic happens in the close() method of the external WebSocket communication module when closing one of the tracked peer streams fails. Each stream is locked, and stream.close().await is expected to succeed; a failure means the WebSocket transport could not perform the close handshake (connection already dead, broken underlying socket, or task runtime issue).

Source

Thrown at crates/burn-communication/src/external_comm.rs:133

                max_downloads,
                cur_download_count: 0,
            },
        );
        core::mem::drop(exposed_tensors);
        self.new_tensor_notify.notify_waiters();
    }

    pub async fn close(&self) {
        // Send a closing message to every open WebSocket stream

        let mut streams = self.channels.lock().await;
        for (_, stream) in streams.drain() {
            let mut stream = stream.lock().await;

            stream
                .close()
                .await
                .expect("Failed to close WebSocket stream");
        }
    }

    /// Downloads a tensor that is exposed on another server. Requires a Tokio 1.x runtime
    ///
    /// Returns None if the peer closes the connection
    pub async fn download_tensor(
        &self,
        remote: Address,
        transfer_id: TensorTransferId,
    ) -> Option<TensorData> {
        log::info!("Downloading tensor from {remote:?}");

        let stream = self.get_data_stream(remote).await;
        let mut stream = stream.lock().await;

        // Send the download request with the download id
        let bytes: bytes::Bytes =

View on GitHub (pinned to d16f7ba2ed)

Solutions

Treat close failures as non-fatal: replace expect with a logged warning, since the socket is being discarded anyway.
Check peer liveness and clean stale/dead connections from the streams map before closing.
Ensure the tokio runtime is running when close() is awaited — calling it outside a runtime can error.
Add keepalive/ping timeouts so dead connections are pruned before shutdown.

Example fix

// before
stream.close().await.expect("Failed to close WebSocket stream");
// after
if let Err(e) = stream.close().await {
    log::warn!("failed to close websocket stream gracefully: {e}");
}

Defensive patterns

Strategy: try-catch

Try / catch

if let Err(e) = stream.close().await {
    log::warn!("websocket close failed (ignoring): {e}");
}

Prevention

Treat close errors as non-fatal — the socket is being discarded anyway.
Prune dead connections periodically instead of only at shutdown.
Ensure close() runs inside an active tokio runtime.

When it happens

Trigger: Calling close() (server/communicator shutdown) when a peer connection has already been severed abruptly (peer process died, network drop, TCP reset), so the websocket close handshake errors.

Common situations: Shutting down a tensor-sharing server while remote peers have disconnected; network partitions or firewall kills during long training runs; peer crash leaving half-open sockets that fail on graceful close.

Related errors

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/e581e566018542c6. Report an issue: GitHub.