diff --git a/protocols/kad/src/handler.rs b/protocols/kad/src/handler.rs index 318261d8d21..5e7c2e21b8b 100644 --- a/protocols/kad/src/handler.rs +++ b/protocols/kad/src/handler.rs @@ -60,7 +60,8 @@ pub struct Handler { next_connec_unique_id: UniqueConnecId, /// List of active outbound streams. - outbound_substreams: futures_bounded::FuturesMap>>, + outbound_substreams: + futures_bounded::FuturesTupleSet>, QueryId>, /// Contains one [`oneshot::Sender`] per outbound stream that we have requested. pending_streams: @@ -453,7 +454,7 @@ impl Handler { remote_peer_id, next_connec_unique_id: UniqueConnecId(0), inbound_substreams: Default::default(), - outbound_substreams: futures_bounded::FuturesMap::new( + outbound_substreams: futures_bounded::FuturesTupleSet::new( Duration::from_secs(10), MAX_NUM_STREAMS, ), @@ -552,32 +553,36 @@ impl Handler { let (sender, receiver) = oneshot::channel(); self.pending_streams.push_back(sender); - let result = self.outbound_substreams.try_push(id, async move { - let mut stream = receiver - .await - .map_err(|_| io::Error::from(io::ErrorKind::BrokenPipe))? - .map_err(|e| match e { - StreamUpgradeError::Timeout => io::ErrorKind::TimedOut.into(), - StreamUpgradeError::Apply(e) => e, - StreamUpgradeError::NegotiationFailed => { - io::Error::new(io::ErrorKind::ConnectionRefused, "protocol not supported") - } - StreamUpgradeError::Io(e) => e, - })?; - - let has_answer = !matches!(msg, KadRequestMsg::AddProvider { .. }); - - stream.send(msg).await?; - stream.close().await?; - - if !has_answer { - return Ok(None); - } + let result = self.outbound_substreams.try_push( + async move { + let mut stream = receiver + .await + .map_err(|_| io::Error::from(io::ErrorKind::BrokenPipe))? + .map_err(|e| match e { + StreamUpgradeError::Timeout => io::ErrorKind::TimedOut.into(), + StreamUpgradeError::Apply(e) => e, + StreamUpgradeError::NegotiationFailed => io::Error::new( + io::ErrorKind::ConnectionRefused, + "protocol not supported", + ), + StreamUpgradeError::Io(e) => e, + })?; + + let has_answer = !matches!(msg, KadRequestMsg::AddProvider { .. }); + + stream.send(msg).await?; + stream.close().await?; + + if !has_answer { + return Ok(None); + } - let msg = stream.next().await.ok_or(io::ErrorKind::UnexpectedEof)??; + let msg = stream.next().await.ok_or(io::ErrorKind::UnexpectedEof)??; - Ok(Some(msg)) - }); + Ok(Some(msg)) + }, + id, + ); debug_assert!( result.is_ok(), @@ -728,15 +733,15 @@ impl ConnectionHandler for Handler { } match self.outbound_substreams.poll_unpin(cx) { - Poll::Ready((query, Ok(Ok(Some(response))))) => { + Poll::Ready((Ok(Ok(Some(response))), query_id)) => { return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour( - process_kad_response(response, query), + process_kad_response(response, query_id), )) } - Poll::Ready((_, Ok(Ok(None)))) => { + Poll::Ready((Ok(Ok(None)), _)) => { continue; } - Poll::Ready((query_id, Ok(Err(e)))) => { + Poll::Ready((Ok(Err(e)), query_id)) => { return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour( HandlerEvent::QueryError { error: HandlerQueryErr::Io(e), @@ -744,7 +749,7 @@ impl ConnectionHandler for Handler { }, )) } - Poll::Ready((query_id, Err(_timeout))) => { + Poll::Ready((Err(_timeout), query_id)) => { return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour( HandlerEvent::QueryError { error: HandlerQueryErr::Io(io::ErrorKind::TimedOut.into()),