webdataset_core/
handlers.rs

1//! Pluggable exception handlers.
2//!
3//! These mirror the handlers in the Python implementation. A handler inspects
4//! an [`Error`] and decides what the pipeline stage that produced it should do
5//! next. Because Rust iterators cannot unwind and resume, "re-raise" means
6//! "forward the error to the consumer" rather than "panic".
7//!
8//! ```
9//! use webdataset_core::handlers::{warn_and_continue, Action};
10//! use webdataset_core::Error;
11//!
12//! let handler = warn_and_continue();
13//! assert_eq!(handler.handle(&Error::value("boom")), Action::Continue);
14//! ```
15
16use alloc::sync::Arc;
17#[cfg(feature = "std")]
18use core::time::Duration;
19
20use crate::error::Error;
21
22/// What a pipeline stage should do after an error.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Action {
25    /// Drop the offending item and keep going.
26    Continue,
27    /// Stop producing items; the stream ends cleanly.
28    Stop,
29    /// Forward the error downstream so the consumer sees it.
30    Reraise,
31}
32
33/// Decides how a pipeline stage reacts to an error.
34pub trait Handler: Send + Sync + core::fmt::Debug {
35    /// Inspect `error` and decide what happens next.
36    fn handle(&self, error: &Error) -> Action;
37}
38
39/// A shared, cheaply clonable handler.
40pub type HandlerRef = Arc<dyn Handler>;
41
42macro_rules! simple_handler {
43    ($name:ident, $ctor:ident, $action:expr, $warn:expr, $doc:expr) => {
44        #[doc = $doc]
45        #[derive(Debug, Clone, Copy, Default)]
46        pub struct $name;
47
48        impl Handler for $name {
49            fn handle(&self, error: &Error) -> Action {
50                if $warn {
51                    log::warn!("{error}");
52                    // Matches the Python implementation: slow the stream down a
53                    // little so warnings do not scroll past unnoticed.
54                    #[cfg(feature = "std")]
55                    std::thread::sleep(Duration::from_millis(500));
56                }
57                $action
58            }
59        }
60
61        #[doc = $doc]
62        pub fn $ctor() -> HandlerRef {
63            Arc::new($name)
64        }
65    };
66}
67
68simple_handler!(
69    ReraiseException,
70    reraise_exception,
71    Action::Reraise,
72    false,
73    "Forward the error to the consumer of the pipeline."
74);
75simple_handler!(
76    IgnoreAndContinue,
77    ignore_and_continue,
78    Action::Continue,
79    false,
80    "Silently drop the offending item and continue."
81);
82simple_handler!(
83    WarnAndContinue,
84    warn_and_continue,
85    Action::Continue,
86    true,
87    "Log the error, drop the offending item, and continue."
88);
89simple_handler!(IgnoreAndStop, ignore_and_stop, Action::Stop, false, "Silently end the stream.");
90simple_handler!(WarnAndStop, warn_and_stop, Action::Stop, true, "Log the error and end the stream.");
91
92/// Wraps a closure so it can be used as a [`Handler`].
93pub struct FnHandler<F>(pub F);
94
95impl<F> core::fmt::Debug for FnHandler<F> {
96    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
97        f.write_str("FnHandler")
98    }
99}
100
101impl<F> Handler for FnHandler<F>
102where
103    F: Fn(&Error) -> Action + Send + Sync,
104{
105    fn handle(&self, error: &Error) -> Action {
106        (self.0)(error)
107    }
108}
109
110/// Build a [`HandlerRef`] from a closure.
111pub fn handler_fn<F>(f: F) -> HandlerRef
112where
113    F: Fn(&Error) -> Action + Send + Sync + 'static,
114{
115    Arc::new(FnHandler(f))
116}
117
118/// Apply `handler` to `error` and translate its decision into the action a
119/// pipeline stage should take.
120pub fn dispatch(handler: &dyn Handler, error: Error) -> Dispatch {
121    match handler.handle(&error) {
122        Action::Continue => Dispatch::Skip,
123        Action::Stop => Dispatch::Stop,
124        Action::Reraise => Dispatch::Yield(error),
125    }
126}
127
128/// The outcome of running a [`Handler`] over an error, as consumed by stages.
129#[derive(Debug)]
130pub enum Dispatch {
131    /// Drop the item and pull the next one.
132    Skip,
133    /// End the stream.
134    Stop,
135    /// Yield this error downstream.
136    Yield(Error),
137}
webdataset_core/handlers.rs

webdataset_core/
handlers.rs
