Skip to content

Commit

Permalink
fix(swarm): implement ConnectionHandler::poll_close for combinators
Browse files Browse the repository at this point in the history
Follow-up to #4076.

This is especially relevant since `libp2p-swarm-derive` uses `SelectConnectionHandler`.

Pull-Request: #4794.
  • Loading branch information
mxinden authored Nov 5, 2023
1 parent 32c945c commit 402212c
Show file tree
Hide file tree
Showing 7 changed files with 58 additions and 2 deletions.
8 changes: 8 additions & 0 deletions swarm/src/behaviour/toggle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -363,4 +363,12 @@ where
}
}
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
let Some(inner) = self.inner.as_mut() else {
return Poll::Ready(None);
};

inner.poll_close(cx)
}
}
3 changes: 3 additions & 0 deletions swarm/src/handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,9 @@ pub trait ConnectionHandler: Send + 'static {
/// We therefore cannot guarantee that performing IO within here will succeed.
///
/// To signal completion, [`Poll::Ready(None)`] should be returned.
///
/// Implementations MUST have a [`fuse`](futures::StreamExt::fuse)-like behaviour.
/// That is, [`Poll::Ready(None)`] MUST be returned on repeated calls to [`ConnectionHandler::poll_close`].
fn poll_close(&mut self, _: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
Poll::Ready(None)
}
Expand Down
9 changes: 9 additions & 0 deletions swarm/src/handler/either.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,15 @@ where
Poll::Ready(event)
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
let event = match self {
Either::Left(handler) => futures::ready!(handler.poll_close(cx)).map(Either::Left),
Either::Right(handler) => futures::ready!(handler.poll_close(cx)).map(Either::Right),
};

Poll::Ready(event)
}

fn on_connection_event(
&mut self,
event: ConnectionEvent<
Expand Down
4 changes: 4 additions & 0 deletions swarm/src/handler/map_in.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,10 @@ where
self.inner.poll(cx)
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
self.inner.poll_close(cx)
}

fn on_connection_event(
&mut self,
event: ConnectionEvent<
Expand Down
9 changes: 9 additions & 0 deletions swarm/src/handler/map_out.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
use crate::handler::{
ConnectionEvent, ConnectionHandler, ConnectionHandlerEvent, SubstreamProtocol,
};
use futures::ready;
use std::fmt::Debug;
use std::task::{Context, Poll};

Expand Down Expand Up @@ -83,6 +84,14 @@ where
})
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
let Some(e) = ready!(self.inner.poll_close(cx)) else {
return Poll::Ready(None);
};

Poll::Ready(Some((self.map)(e)))
}

fn on_connection_event(
&mut self,
event: ConnectionEvent<
Expand Down
13 changes: 12 additions & 1 deletion swarm/src/handler/multi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use crate::handler::{
};
use crate::upgrade::{InboundUpgradeSend, OutboundUpgradeSend, UpgradeInfoSend};
use crate::Stream;
use futures::{future::BoxFuture, prelude::*};
use futures::{future::BoxFuture, prelude::*, ready};
use rand::Rng;
use std::{
cmp,
Expand Down Expand Up @@ -271,6 +271,17 @@ where

Poll::Pending
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
for (k, h) in self.handlers.iter_mut() {
let Some(e) = ready!(h.poll_close(cx)) else {
continue;
};
return Poll::Ready(Some((k.clone(), e)));
}

Poll::Ready(None)
}
}

/// Split [`MultiHandler`] into parts.
Expand Down
14 changes: 13 additions & 1 deletion swarm/src/handler/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use crate::handler::{
};
use crate::upgrade::SendWrapper;
use either::Either;
use futures::future;
use futures::{future, ready};
use libp2p_core::upgrade::SelectUpgrade;
use std::{cmp, task::Context, task::Poll};

Expand Down Expand Up @@ -259,6 +259,18 @@ where
Poll::Pending
}

fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
if let Some(e) = ready!(self.proto1.poll_close(cx)) {
return Poll::Ready(Some(Either::Left(e)));
}

if let Some(e) = ready!(self.proto2.poll_close(cx)) {
return Poll::Ready(Some(Either::Right(e)));
}

Poll::Ready(None)
}

fn on_connection_event(
&mut self,
event: ConnectionEvent<
Expand Down

0 comments on commit 402212c

Please sign in to comment.