feat(core): add post-round trace action to run_with and spawn_with callbacks`

Author: fujiapple852

crates/trippy-core/src/lib.rs

@@ -82,7 +82,7 @@ pub use probe::{
     ProbeComplete, ProbeStatus, UnknownExtension,
 };
 pub use state::{Hop, NatStatus, State};
-pub use strategy::{CompletionReason, Round, Strategy};
+pub use strategy::{Action, CompletionReason, Round, Strategy};
 pub use tracer::Tracer;
 pub use types::{
     Dscp, Ecn, Flags, MaxInflight, MaxRounds, PacketSize, PayloadPattern, Port, RoundId, Sequence,

crates/trippy-core/src/strategy.rs

@@ -49,14 +49,29 @@ pub enum CompletionReason {
     RoundTimeLimitExceeded,
 }
 
+/// The action to take after the completion of a round.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum Action {
+    /// Continue tracing.
+    Continue,
+    /// Stop tracing.
+    Stop,
+}
+
+impl From<()> for Action {
+    fn from((): ()) -> Self {
+        Self::Continue
+    }
+}
+
 /// Trace a path to a target.
 #[derive(Debug, Clone)]
 pub struct Strategy<F> {
     config: StrategyConfig,
     publish: F,
 }
 
-impl<F: Fn(&Round<'_>)> Strategy<F> {
+impl<F: Fn(&Round<'_>) -> Action> Strategy<F> {
     #[instrument(skip_all, level = "trace")]
     pub fn new(config: &StrategyConfig, publish: F) -> Self {
         tracing::debug!(?config);
@@ -192,8 +207,8 @@ impl<F: Fn(&Round<'_>)> Strategy<F> {
         let round_max = round_duration > self.config.max_round_duration;
         let target_found = st.target_found();
         if round_min && grace_exceeded && target_found || round_max {
-            self.publish_trace(st);
-            st.advance_round(self.config.first_ttl);
+            let action = self.publish_trace(st);
+            st.advance_round(self.config.first_ttl, action);
         }
     }
 
@@ -202,7 +217,7 @@ impl<F: Fn(&Round<'_>)> Strategy<F> {
     /// If the round completed without receiving an `EchoReply` from the target host then we also
     /// publish the next `ProbeStatus` which is assumed to represent the TTL of the target host.
     #[instrument(skip(self, state), level = "trace")]
-    fn publish_trace(&self, state: &TracerState) {
+    fn publish_trace(&self, state: &TracerState) -> Action {
         let max_received_ttl = if let Some(target_ttl) = state.target_ttl() {
             target_ttl
         } else {
@@ -220,7 +235,7 @@ impl<F: Fn(&Round<'_>)> Strategy<F> {
         } else {
             CompletionReason::RoundTimeLimitExceeded
         };
-        (self.publish)(&Round::new(probes, largest_ttl, reason));
+        (self.publish)(&Round::new(probes, largest_ttl, reason))
     }
 
     /// Check if the `TraceId` matches the expected value for this tracer.
@@ -840,7 +855,7 @@ mod tests {
             protocol: Protocol::Tcp,
             ..Default::default()
         };
-        let tracer = Strategy::new(&config, |_| {});
+        let tracer = Strategy::new(&config, |_| Action::Continue);
         let mut state = TracerState::new(config);
         tracer.send_request(&mut network, &mut state)?;
         tracer.recv_response(&mut network, &mut state)?;
@@ -870,7 +885,7 @@ mod state {
     use crate::probe::{Probe, ProbeStatus};
     use crate::strategy::{StrategyConfig, StrategyResponse};
     use crate::types::{MaxRounds, Port, RoundId, Sequence, TimeToLive, TraceId};
-    use crate::{Flags, MultipathStrategy, PortDirection, Protocol};
+    use crate::{Action, Flags, MultipathStrategy, PortDirection, Protocol};
     use std::array::from_fn;
     use std::net::IpAddr;
     use std::time::SystemTime;
@@ -929,6 +944,8 @@ mod state {
         target_ttl: Option<TimeToLive>,
         /// The timestamp of the echo response packet.
         received_time: Option<SystemTime>,
+        /// The action to take before starting the next round.
+        next_round_action: Action,
     }
 
     impl TracerState {
@@ -945,6 +962,7 @@ mod state {
                 max_received_ttl: None,
                 target_ttl: None,
                 received_time: None,
+                next_round_action: Action::Continue,
             }
         }
 
@@ -995,11 +1013,12 @@ mod state {
         }
 
         /// Are all rounds complete?
-        pub const fn finished(&self, max_rounds: Option<MaxRounds>) -> bool {
-            match max_rounds {
-                None => false,
-                Some(max_rounds) => self.round.0 > max_rounds.0.get() - 1,
-            }
+        pub fn finished(&self, max_rounds: Option<MaxRounds>) -> bool {
+            self.next_round_action == Action::Stop
+                || match max_rounds {
+                    None => false,
+                    Some(max_rounds) => self.round.0 > max_rounds.0.get() - 1,
+                }
         }
 
         /// Create and return the next `Probe` at the current `sequence` and `ttl`.
@@ -1257,7 +1276,8 @@ mod state {
         /// reset it here. We do this here to avoid having to deal with the sequence number
         /// wrapping during a round, which is more problematic.
         #[instrument(skip(self), level = "trace")]
-        pub fn advance_round(&mut self, first_ttl: TimeToLive) {
+        pub fn advance_round(&mut self, first_ttl: TimeToLive, next_round_action: Action) {
+            self.next_round_action = next_round_action;
             if self.sequence >= self.max_sequence() {
                 self.sequence = self.config.initial_sequence;
             }
@@ -1375,7 +1395,7 @@ mod state {
             }
 
             // Advance to the next round
-            state.advance_round(TimeToLive(1));
+            state.advance_round(TimeToLive(1), Action::Continue);
 
             // Validate the `TracerState` after the round update
             assert_eq!(state.round, RoundId(1));
@@ -1509,7 +1529,7 @@ mod state {
             }
 
             // Advance the round, which will wrap the sequence back to `initial_sequence`
-            state.advance_round(TimeToLive(1));
+            state.advance_round(TimeToLive(1), Action::Continue);
             assert_eq!(state.round, RoundId(1));
             assert_eq!(state.sequence, initial_sequence);
             assert_eq!(state.round_sequence, initial_sequence);
@@ -1547,7 +1567,7 @@ mod state {
                 for _ in 0..max_probe_per_round {
                     let _probe = state.next_probe(SystemTime::now());
                 }
-                state.advance_round(TimeToLive(1));
+                state.advance_round(TimeToLive(1), Action::Continue);
             }
             assert_eq!(state.round, RoundId(2000));
             assert_eq!(state.round_sequence, Sequence(33434));
@@ -1564,7 +1584,7 @@ mod state {
                 for _ in 0..rng.random_range(0..max_probe_per_round) {
                     state.next_probe(SystemTime::now());
                 }
-                state.advance_round(TimeToLive(1));
+                state.advance_round(TimeToLive(1), Action::Continue);
             }
         }
 
@@ -1578,7 +1598,7 @@ mod state {
                     _ = state.next_probe(SystemTime::now());
                     _ = state.reissue_probe(SystemTime::now());
                 }
-                state.advance_round(TimeToLive(1));
+                state.advance_round(TimeToLive(1), Action::Continue);
             }
             assert_eq!(state.round, RoundId(2000));
             assert_eq!(state.round_sequence, Sequence(57310));
@@ -1600,7 +1620,7 @@ mod state {
             for _ in 0..55 {
                 _ = state.next_probe(SystemTime::now());
             }
-            state.advance_round(TimeToLive(1));
+            state.advance_round(TimeToLive(1), Action::Continue);
             assert!(!state.in_round(Sequence(64491)));
         }
 

crates/trippy-core/src/tracer.rs

@@ -1,6 +1,6 @@
 use crate::error::Result;
 use crate::{
-    Error, IcmpExtensionParseMode, MaxInflight, MaxRounds, MultipathStrategy, PacketSize,
+    Action, Error, IcmpExtensionParseMode, MaxInflight, MaxRounds, MultipathStrategy, PacketSize,
     PayloadPattern, PortDirection, PrivilegeMode, Protocol, Round, Sequence, State, TimeToLive,
     TraceId, TypeOfService,
 };
@@ -137,9 +137,12 @@ impl Tracer {
     /// retrieved using the [`Tracer::snapshot`] method.
     ///
     /// This method will additionally call the provided function for each round
-    /// that is completed.  This can be useful if you want to gather round state
+    /// that is completed. This can be useful if you want to gather round state
     /// manually if the tracer is run indefinitely (by not setting
-    /// [`crate::Builder::max_rounds`])
+    /// [`crate::Builder::max_rounds`]).
+    ///
+    /// The callback may either return `()` to continue tracing, or return
+    /// [`Action`] to decide whether tracing should continue after each round.
     ///
     /// # Example
     ///
@@ -159,10 +162,39 @@ impl Tracer {
     /// # }
     /// ```
     ///
+    /// The following will stop after the first round for which `stop` is set:
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # fn main() -> anyhow::Result<()> {
+    /// # use std::net::IpAddr;
+    /// # use std::str::FromStr;
+    /// # use std::sync::atomic::{AtomicBool, Ordering};
+    /// use trippy_core::{Builder, Action};
+    ///
+    /// let addr = IpAddr::from_str("1.1.1.1")?;
+    /// let stop = AtomicBool::new(false);
+    /// let tracer = Builder::new(addr).build()?;
+    /// tracer.run_with(|_| {
+    ///     if stop.load(Ordering::Relaxed) {
+    ///         Action::Stop
+    ///     } else {
+    ///         Action::Continue
+    ///     }
+    /// })?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
     /// # See Also
     ///
     /// - [`Tracer::run`] - Run the tracer without a custom round handler.
-    pub fn run_with<F: Fn(&Round<'_>)>(&self, func: F) -> Result<()> {
+    pub fn run_with<F, R>(&self, func: F) -> Result<()>
+    where
+        F: Fn(&Round<'_>) -> R,
+        R: Into<Action>,
+    {
         self.inner.run_with(func)
     }
 
@@ -249,13 +281,47 @@ impl Tracer {
     /// # }
     /// ```
     ///
+    /// The callback may either return `()` to continue tracing, or return
+    /// [`Action`] to decide whether tracing should continue after each round.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # fn main() -> anyhow::Result<()> {
+    /// # use std::net::IpAddr;
+    /// # use std::str::FromStr;
+    /// # use std::sync::{
+    /// #     Arc,
+    /// #     atomic::{AtomicBool, Ordering},
+    /// # };
+    /// use trippy_core::{Builder, Action};
+    ///
+    /// let addr = IpAddr::from_str("1.1.1.1")?;
+    /// let stop = Arc::new(AtomicBool::new(false));
+    /// let stop_for_trace = Arc::clone(&stop);
+    /// let (tracer, handle) = Builder::new(addr)
+    ///     .build()?
+    ///     .spawn_with(move |_| {
+    ///         if stop_for_trace.load(Ordering::Relaxed) {
+    ///             Action::Stop
+    ///         } else {
+    ///             Action::Continue
+    ///         }
+    ///     })?;
+    /// stop.store(true, Ordering::Relaxed);
+    /// handle.join().unwrap()?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
     /// # See Also
     ///
     /// - [`Tracer::spawn`] - Spawn the tracer on a new thread without a custom round handler.
-    pub fn spawn_with<F: Fn(&Round<'_>) + Send + 'static>(
-        self,
-        func: F,
-    ) -> Result<(Self, JoinHandle<Result<()>>)> {
+    pub fn spawn_with<F, R>(self, func: F) -> Result<(Self, JoinHandle<Result<()>>)>
+    where
+        F: Fn(&Round<'_>) -> R + Send + 'static,
+        R: Into<Action> + Send + 'static,
+    {
         let tracer = self.clone();
         let handle = thread::Builder::new()
             .name(format!("tracer-{}", self.trace_identifier().0))
@@ -425,7 +491,7 @@ mod inner {
     use crate::error::Result;
     use crate::net::{PlatformImpl, SocketImpl};
     use crate::{
-        Channel, Error, IcmpExtensionParseMode, MaxInflight, MaxRounds, MultipathStrategy,
+        Action, Channel, Error, IcmpExtensionParseMode, MaxInflight, MaxRounds, MultipathStrategy,
         PacketSize, PayloadPattern, PortDirection, PrivilegeMode, Protocol, Round, Sequence,
         SourceAddr, State, Strategy, TimeToLive, TraceId, TypeOfService,
     };
@@ -530,12 +596,15 @@ mod inner {
 
         #[instrument(skip_all, level = "trace")]
         pub(super) fn run(&self) -> Result<()> {
-            self.run_internal(|_| ())
-                .map_err(|err| self.handle_error(err))
+            self.run_with(|_| ())
         }
 
         #[instrument(skip_all, level = "trace")]
-        pub(super) fn run_with<F: Fn(&Round<'_>)>(&self, func: F) -> Result<()> {
+        pub(super) fn run_with<F, R>(&self, func: F) -> Result<()>
+        where
+            F: Fn(&Round<'_>) -> R,
+            R: Into<Action>,
+        {
             self.run_internal(func)
                 .map_err(|err| self.handle_error(err))
         }
@@ -646,7 +715,11 @@ mod inner {
         }
 
         #[instrument(skip_all, level = "trace")]
-        fn run_internal<F: Fn(&Round<'_>)>(&self, func: F) -> Result<()> {
+        fn run_internal<F, R>(&self, func: F) -> Result<()>
+        where
+            F: Fn(&Round<'_>) -> R,
+            R: Into<Action>,
+        {
             // if we are given a source address, validate it otherwise
             // discover it based on the target address and interface.
             let source_addr = match self.source_addr {
@@ -668,7 +741,7 @@ mod inner {
             let strategy_config = self.make_strategy_config();
             let strategy = Strategy::new(&strategy_config, |round| {
                 self.handler(round);
-                func(round);
+                func(round).into()
             });
             strategy.run(channel)?;
             Ok(())