Skip to content

Commit af7a567

Browse files
committed
feat: add support for using poll instead of select and expose socket-readiness-mode config (#1790)
1 parent 76a76c0 commit af7a567

23 files changed

Lines changed: 252 additions & 60 deletions

crates/trippy-core/src/builder.rs

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use crate::constants::MAX_INITIAL_SEQUENCE;
33
use crate::error::Result;
44
use crate::{
55
Error, IcmpExtensionParseMode, MAX_TTL, MaxInflight, MaxRounds, MultipathStrategy, PacketSize,
6-
PayloadPattern, PortDirection, PrivilegeMode, Protocol, Sequence, TimeToLive, TraceId, Tracer,
7-
TypeOfService,
6+
PayloadPattern, PortDirection, PrivilegeMode, Protocol, Sequence, SocketReadinessMode,
7+
TimeToLive, TraceId, Tracer, TypeOfService,
88
};
99
use std::net::IpAddr;
1010
use std::num::NonZeroUsize;
@@ -47,6 +47,7 @@ pub struct Builder {
4747
tos: TypeOfService,
4848
icmp_extension_parse_mode: IcmpExtensionParseMode,
4949
read_timeout: Duration,
50+
socket_readiness_mode: SocketReadinessMode,
5051
tcp_connect_timeout: Duration,
5152
trace_identifier: TraceId,
5253
max_rounds: Option<MaxRounds>,
@@ -77,6 +78,7 @@ impl Default for Builder {
7778
tos: ChannelConfig::default().tos,
7879
icmp_extension_parse_mode: ChannelConfig::default().icmp_extension_parse_mode,
7980
read_timeout: ChannelConfig::default().read_timeout,
81+
socket_readiness_mode: ChannelConfig::default().socket_readiness_mode,
8082
tcp_connect_timeout: ChannelConfig::default().tcp_connect_timeout,
8183
trace_identifier: StrategyConfig::default().trace_identifier,
8284
max_rounds: StrategyConfig::default().max_rounds,
@@ -381,6 +383,31 @@ impl Builder {
381383
}
382384
}
383385

386+
/// Set the socket readiness backend.
387+
///
388+
/// The default is [`SocketReadinessMode::Select`].
389+
///
390+
/// # Examples
391+
///
392+
/// ```no_run
393+
/// # fn main() -> anyhow::Result<()> {
394+
/// use trippy_core::{Builder, SocketReadinessMode};
395+
///
396+
/// let addr = std::net::IpAddr::from([1, 1, 1, 1]);
397+
/// let tracer = Builder::new(addr)
398+
/// .socket_readiness_mode(SocketReadinessMode::Poll)
399+
/// .build()?;
400+
/// # Ok(())
401+
/// # }
402+
/// ```
403+
#[must_use]
404+
pub fn socket_readiness_mode(self, socket_readiness_mode: SocketReadinessMode) -> Self {
405+
Self {
406+
socket_readiness_mode,
407+
..self
408+
}
409+
}
410+
384411
/// Set the TCP connect timeout.
385412
///
386413
/// # Examples
@@ -743,6 +770,7 @@ impl Builder {
743770
self.tos,
744771
self.icmp_extension_parse_mode,
745772
self.read_timeout,
773+
self.socket_readiness_mode,
746774
self.tcp_connect_timeout,
747775
self.trace_identifier,
748776
self.max_rounds,
@@ -805,6 +833,10 @@ mod tests {
805833
defaults::DEFAULT_STRATEGY_READ_TIMEOUT,
806834
tracer.read_timeout()
807835
);
836+
assert_eq!(
837+
defaults::DEFAULT_SOCKET_READINESS_MODE,
838+
tracer.socket_readiness_mode()
839+
);
808840
assert_eq!(
809841
defaults::DEFAULT_STRATEGY_TCP_CONNECT_TIMEOUT,
810842
tracer.tcp_connect_timeout()
@@ -851,6 +883,7 @@ mod tests {
851883
.tos(0x1a)
852884
.icmp_extension_parse_mode(IcmpExtensionParseMode::Enabled)
853885
.read_timeout(Duration::from_millis(50))
886+
.socket_readiness_mode(SocketReadinessMode::Poll)
854887
.tcp_connect_timeout(Duration::from_millis(100))
855888
.max_rounds(Some(10))
856889
.first_ttl(2)
@@ -882,6 +915,7 @@ mod tests {
882915
tracer.icmp_extension_parse_mode()
883916
);
884917
assert_eq!(Duration::from_millis(50), tracer.read_timeout());
918+
assert_eq!(SocketReadinessMode::Poll, tracer.socket_readiness_mode());
885919
assert_eq!(Duration::from_millis(100), tracer.tcp_connect_timeout());
886920
assert_eq!(
887921
Some(MaxRounds(NonZeroUsize::new(10).unwrap())),

crates/trippy-core/src/config.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use std::time::Duration;
99

1010
/// Default values for configuration.
1111
pub mod defaults {
12-
use crate::config::IcmpExtensionParseMode;
12+
use crate::config::{IcmpExtensionParseMode, SocketReadinessMode};
1313
use crate::{MultipathStrategy, PrivilegeMode, Protocol};
1414
use std::time::Duration;
1515

@@ -56,6 +56,9 @@ pub mod defaults {
5656
/// The default value for `read-timeout`.
5757
pub const DEFAULT_STRATEGY_READ_TIMEOUT: Duration = Duration::from_millis(10);
5858

59+
/// The default value for `socket-readiness-mode`.
60+
pub const DEFAULT_SOCKET_READINESS_MODE: SocketReadinessMode = SocketReadinessMode::Select;
61+
5962
/// The default value for `grace-duration`.
6063
pub const DEFAULT_STRATEGY_GRACE_DURATION: Duration = Duration::from_millis(100);
6164

@@ -125,6 +128,24 @@ impl Display for IcmpExtensionParseMode {
125128
}
126129
}
127130

131+
/// The socket readiness backend.
132+
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
133+
pub enum SocketReadinessMode {
134+
/// Use `select` for socket readiness checks.
135+
Select,
136+
/// Use `poll` for socket readiness checks.
137+
Poll,
138+
}
139+
140+
impl Display for SocketReadinessMode {
141+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142+
match self {
143+
Self::Select => write!(f, "select"),
144+
Self::Poll => write!(f, "poll"),
145+
}
146+
}
147+
}
148+
128149
/// The tracing protocol.
129150
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
130151
pub enum Protocol {
@@ -282,6 +303,7 @@ pub struct ChannelConfig {
282303
pub tos: TypeOfService,
283304
pub icmp_extension_parse_mode: IcmpExtensionParseMode,
284305
pub read_timeout: Duration,
306+
pub socket_readiness_mode: SocketReadinessMode,
285307
pub tcp_connect_timeout: Duration,
286308
}
287309

@@ -298,6 +320,7 @@ impl Default for ChannelConfig {
298320
tos: TypeOfService(defaults::DEFAULT_STRATEGY_TOS),
299321
icmp_extension_parse_mode: defaults::DEFAULT_ICMP_EXTENSION_PARSE_MODE,
300322
read_timeout: defaults::DEFAULT_STRATEGY_READ_TIMEOUT,
323+
socket_readiness_mode: defaults::DEFAULT_SOCKET_READINESS_MODE,
301324
tcp_connect_timeout: defaults::DEFAULT_STRATEGY_TCP_CONNECT_TIMEOUT,
302325
}
303326
}

crates/trippy-core/src/error.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ pub enum IoOperation {
7979
NewSocket,
8080
SetNonBlocking,
8181
Select,
82+
Poll,
8283
RecvFrom,
8384
Read,
8485
Shutdown,
@@ -110,6 +111,7 @@ impl Display for IoOperation {
110111
Self::NewSocket => write!(f, "create new socket"),
111112
Self::SetNonBlocking => write!(f, "set non-blocking"),
112113
Self::Select => write!(f, "select"),
114+
Self::Poll => write!(f, "poll"),
113115
Self::RecvFrom => write!(f, "recv from"),
114116
Self::Read => write!(f, "read"),
115117
Self::Shutdown => write!(f, "shutdown"),

crates/trippy-core/src/lib.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,8 @@ use net::source::SourceAddr;
7272

7373
pub use builder::Builder;
7474
pub use config::{
75-
IcmpExtensionParseMode, MultipathStrategy, PortDirection, PrivilegeMode, Protocol, defaults,
75+
IcmpExtensionParseMode, MultipathStrategy, PortDirection, PrivilegeMode, Protocol,
76+
SocketReadinessMode, defaults,
7677
};
7778
pub use constants::MAX_TTL;
7879
pub use error::Error;

crates/trippy-core/src/net/channel.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::error::{Error, Result};
33
use crate::net::socket::Socket;
44
use crate::net::{Network, ipv4::Ipv4, ipv6::Ipv6, platform};
55
use crate::probe::{Probe, Response};
6-
use crate::{Port, PrivilegeMode, Protocol};
6+
use crate::{Port, PrivilegeMode, Protocol, SocketReadinessMode};
77
use arrayvec::ArrayVec;
88
use std::net::IpAddr;
99
use std::time::{Duration, SystemTime};
@@ -19,6 +19,7 @@ const MAX_TCP_PROBES: usize = 256;
1919
pub struct Channel<S: Socket> {
2020
protocol: Protocol,
2121
read_timeout: Duration,
22+
socket_readiness_mode: SocketReadinessMode,
2223
tcp_connect_timeout: Duration,
2324
send_socket: Option<S>,
2425
recv_socket: S,
@@ -79,6 +80,7 @@ impl<S: Socket> Channel<S> {
7980
Ok(Self {
8081
protocol: config.protocol,
8182
read_timeout: config.read_timeout,
83+
socket_readiness_mode: config.socket_readiness_mode,
8284
tcp_connect_timeout: config.tcp_connect_timeout,
8385
send_socket,
8486
recv_socket,
@@ -154,7 +156,10 @@ impl<S: Socket> Channel<S> {
154156
/// Generate a `ProbeResponse` for the next available ICMP packet, if any
155157
#[instrument(skip(self), level = "trace")]
156158
fn recv_icmp_probe(&mut self) -> Result<Option<Response>> {
157-
if self.recv_socket.is_readable(self.read_timeout)? {
159+
if self
160+
.recv_socket
161+
.is_readable(self.read_timeout, self.socket_readiness_mode)?
162+
{
158163
match &self.family_config {
159164
FamilyConfig::V4(ipv4) => ipv4.recv_icmp_probe(&mut self.recv_socket),
160165
FamilyConfig::V6(ipv6) => ipv6.recv_icmp_probe(&mut self.recv_socket),
@@ -177,7 +182,11 @@ impl<S: Socket> Channel<S> {
177182
.iter_mut()
178183
.enumerate()
179184
.find_map(|(index, probe)| {
180-
if probe.socket.is_writable().unwrap_or_default() {
185+
if probe
186+
.socket
187+
.is_writable(self.socket_readiness_mode)
188+
.unwrap_or_default()
189+
{
181190
Some(index)
182191
} else {
183192
None

crates/trippy-core/src/net/platform/unix.rs

Lines changed: 70 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -150,12 +150,14 @@ mod address {
150150
}
151151

152152
mod socket {
153+
use crate::SocketReadinessMode;
153154
use crate::error::{ErrorKind, IoError, IoOperation};
154155
use crate::error::{IoResult, Result};
155156
use crate::net::socket::{Socket, SocketError};
156157
use itertools::Itertools;
157158
use nix::{
158159
Error,
160+
poll::{PollFd, PollFlags, PollTimeout, poll},
159161
sys::select::FdSet,
160162
sys::time::{TimeVal, TimeValLike},
161163
};
@@ -227,6 +229,66 @@ mod socket {
227229
.map_err(|err| IoError::Other(err, IoOperation::LocalAddr))?
228230
.as_socket())
229231
}
232+
233+
fn is_readable_select(&self, timeout: Duration) -> IoResult<bool> {
234+
let mut read = FdSet::new();
235+
read.insert(self.inner.as_fd());
236+
let readable = nix::sys::select::select(
237+
None,
238+
Some(&mut read),
239+
None,
240+
None,
241+
Some(&mut TimeVal::milliseconds(timeout.as_millis() as i64)),
242+
);
243+
match readable {
244+
Ok(readable) => Ok(readable == 1),
245+
Err(Error::EINTR) => Ok(false),
246+
Err(err) => Err(IoError::Other(io::Error::from(err), IoOperation::Select)),
247+
}
248+
}
249+
250+
fn is_writable_select(&self) -> IoResult<bool> {
251+
let mut write = FdSet::new();
252+
write.insert(self.inner.as_fd());
253+
let writable = nix::sys::select::select(
254+
None,
255+
None,
256+
Some(&mut write),
257+
None,
258+
Some(&mut TimeVal::zero()),
259+
);
260+
match writable {
261+
Ok(writable) => Ok(writable == 1),
262+
Err(Error::EINTR) => Ok(false),
263+
Err(err) => Err(IoError::Other(io::Error::from(err), IoOperation::Select)),
264+
}
265+
}
266+
267+
fn is_readable_poll(&self, timeout: Duration) -> IoResult<bool> {
268+
let timeout = PollTimeout::try_from(timeout).map_err(|err| {
269+
IoError::Other(
270+
io::Error::new(io::ErrorKind::InvalidInput, err),
271+
IoOperation::Poll,
272+
)
273+
})?;
274+
self.is_ready_poll(PollFlags::POLLIN, timeout)
275+
}
276+
277+
fn is_writable_poll(&self) -> IoResult<bool> {
278+
self.is_ready_poll(PollFlags::POLLOUT, PollTimeout::ZERO)
279+
}
280+
281+
fn is_ready_poll(&self, events: PollFlags, timeout: PollTimeout) -> IoResult<bool> {
282+
let mut fds = [PollFd::new(self.inner.as_fd(), events)];
283+
match poll(&mut fds, timeout) {
284+
Ok(ready) => Ok(ready == 1
285+
&& fds[0].revents().is_some_and(|revents| {
286+
revents.intersects(events | PollFlags::POLLERR | PollFlags::POLLHUP)
287+
})),
288+
Err(Error::EINTR) => Ok(false),
289+
Err(err) => Err(IoError::Other(io::Error::from(err), IoOperation::Poll)),
290+
}
291+
}
230292
}
231293

232294
impl Socket for SocketImpl {
@@ -386,43 +448,17 @@ mod socket {
386448
Ok(())
387449
}
388450
#[instrument(skip(self), level = "trace")]
389-
fn is_readable(&mut self, timeout: Duration) -> IoResult<bool> {
390-
let mut read = FdSet::new();
391-
read.insert(self.inner.as_fd());
392-
let readable = nix::sys::select::select(
393-
None,
394-
Some(&mut read),
395-
None,
396-
None,
397-
Some(&mut TimeVal::milliseconds(timeout.as_millis() as i64)),
398-
);
399-
match readable {
400-
Ok(readable) => Ok(readable == 1),
401-
Err(Error::EINTR) => Ok(false),
402-
Err(err) => Err(IoError::Other(
403-
std::io::Error::from(err),
404-
IoOperation::Select,
405-
)),
451+
fn is_readable(&mut self, timeout: Duration, mode: SocketReadinessMode) -> IoResult<bool> {
452+
match mode {
453+
SocketReadinessMode::Select => self.is_readable_select(timeout),
454+
SocketReadinessMode::Poll => self.is_readable_poll(timeout),
406455
}
407456
}
408457
#[instrument(skip(self), level = "trace")]
409-
fn is_writable(&mut self) -> IoResult<bool> {
410-
let mut write = FdSet::new();
411-
write.insert(self.inner.as_fd());
412-
let writable = nix::sys::select::select(
413-
None,
414-
None,
415-
Some(&mut write),
416-
None,
417-
Some(&mut TimeVal::zero()),
418-
);
419-
match writable {
420-
Ok(writable) => Ok(writable == 1),
421-
Err(Error::EINTR) => Ok(false),
422-
Err(err) => Err(IoError::Other(
423-
std::io::Error::from(err),
424-
IoOperation::Select,
425-
)),
458+
fn is_writable(&mut self, mode: SocketReadinessMode) -> IoResult<bool> {
459+
match mode {
460+
SocketReadinessMode::Select => self.is_writable_select(),
461+
SocketReadinessMode::Poll => self.is_writable_poll(),
426462
}
427463
}
428464
#[instrument(skip(self, buf), level = "trace")]

crates/trippy-core/src/net/platform/windows.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use super::byte_order::Ipv4ByteOrder;
2+
use crate::SocketReadinessMode;
23
use crate::error::{Error, ErrorKind, IoError, IoOperation, IoResult, Result};
34
use crate::net::channel::MAX_PACKET_SIZE;
45
use crate::net::platform::Platform;
@@ -483,7 +484,7 @@ impl Socket for SocketImpl {
483484
}
484485

485486
#[instrument(skip(self), level = "trace")]
486-
fn is_readable(&mut self, timeout: Duration) -> IoResult<bool> {
487+
fn is_readable(&mut self, timeout: Duration, _mode: SocketReadinessMode) -> IoResult<bool> {
487488
if !self.wait_for_event(timeout)? {
488489
return Ok(false);
489490
}
@@ -498,7 +499,7 @@ impl Socket for SocketImpl {
498499
}
499500

500501
#[instrument(skip(self), level = "trace")]
501-
fn is_writable(&mut self) -> IoResult<bool> {
502+
fn is_writable(&mut self, _mode: SocketReadinessMode) -> IoResult<bool> {
502503
if !self.wait_for_event(Duration::ZERO)? {
503504
return Ok(false);
504505
}

0 commit comments

Comments
 (0)