diff --git a/src/multistream_select/dialer_select.rs b/src/multistream_select/dialer_select.rs index 0bd9c259..e3d48230 100644 --- a/src/multistream_select/dialer_select.rs +++ b/src/multistream_select/dialer_select.rs @@ -548,9 +548,7 @@ mod tests { io.close().await.unwrap(); }); - // TODO: Once https://github.com/paritytech/litep2p/pull/62 is merged, this - // should be changed to `is_ok`. - assert!(tokio::time::timeout(Duration::from_secs(10), client).await.is_err()); + assert!(tokio::time::timeout(Duration::from_secs(10), client).await.is_ok()); } #[tokio::test] diff --git a/src/multistream_select/negotiated.rs b/src/multistream_select/negotiated.rs index 2a08f29c..ee701e74 100644 --- a/src/multistream_select/negotiated.rs +++ b/src/multistream_select/negotiated.rs @@ -323,15 +323,22 @@ where } fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Ensure all data has been flushed and expected negotiation messages - // have been received. - ready!(self.as_mut().poll(cx).map_err(Into::::into)?); + // Ensure all data has been flushed, including optimistic multistream-select messages. ready!(self.as_mut().poll_flush(cx).map_err(Into::::into)?); // Continue with the shutdown of the underlying I/O stream. match self.project().state.project() { StateProj::Completed { io, .. } => io.poll_close(cx), - StateProj::Expecting { io, .. } => io.poll_close(cx), + StateProj::Expecting { io, .. } => { + let close_poll = io.poll_close(cx); + if let Poll::Ready(Ok(())) = close_poll { + tracing::debug!( + target: LOG_TARGET, + "Stream closed. Confirmation from remote for optimstic protocol negotiation still pending." + ); + } + close_poll + } StateProj::Invalid => panic!("Negotiated: Invalid state"), } }