1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
use crate::{ErrorHandler, RetryPolicy};
use futures::{Async, AsyncSink, Future, Poll, Sink, StartSend};
use std::time::Instant;
use tokio_timer;

/// Provides a way to handle errors during a `Sink` execution, i.e. it gives you an ability to
/// flush item's with delay.
///
/// This type is similar to [`StreamRetry`](struct.StreamRetry.html). The diffrence is that
/// SinkRetry is more limited:
/// * SinkItem has to implement `Clone` trait.
/// * error_action `OutError` has to implement `From<SinkError>` trait.
///
/// A `fs-readwrite` example is available in the `examples` folder.
/// A `tcp-copy` example handle StreamRetry and SinkRetry example at once.
///
/// A typical usage might be recovering from errors when flushing data to I/O.
///
/// # Warning (Implementation details)
///
/// It depends from inner sink what happen when inner `start_send()` is resolve to error. This
/// function assume that error has the same meaning as `AsyncSink::NotReady`. If your item will
/// be buffered before error was returned - SinkRetry will insert this item again anyway.
pub struct SinkRetry<F, S> {
    error_action: F,
    sink: S,
    state: RetryState,
}

impl<F, S> SinkRetry<F, S>
where
    S: Sink,
    S::SinkItem: Clone,
    F: ErrorHandler<S::SinkError>,
    F::OutError: From<S::SinkError>,
{
    /// Creates a `SinkRetry` using a provided sink and an object of `ErrorHandler` type that
    /// decides on a retry-policy depending on an encountered error.
    ///
    /// # Arguments
    ///
    /// * `sink`: a sink to be filled,
    /// * `error_action`: a type that handles an error and decides which route to take: simply
    ///                   try again, wait and then try, or give up (on a critical error for
    ///                   exapmle).
    /// # Notes
    /// More common use is in like chain manner. See [SinkRetryExt](trait.SinkRetryExt.html)
    pub fn new(sink: S, error_action: F) -> Self
    where
        S: Sink,
    {
        Self {
            error_action,
            sink,
            state: RetryState::WaitingForSink,
        }
    }

    fn try_send_item(&mut self, item: S::SinkItem) -> StartSend<S::SinkItem, F::OutError> {
        debug_assert!(self.state.is_waiting_for_sink());
        loop {
            let cloned_item = item.clone();
            match self.sink.start_send(cloned_item) {
                Err(err) => {
                    match self.error_action.handle(err) {
                        RetryPolicy::Repeat => continue,
                        RetryPolicy::WaitRetry(duration) => {
                            // before return AsyncSink::NotReady(item) we HAVE TO call timer poll
                            // function.
                            let mut timer = tokio_timer::Delay::new(Instant::now() + duration);

                            match timer.poll().expect("Timer panic!") {
                                Async::Ready(_) => match self.poll_complete()? {
                                    Async::Ready(()) => continue,
                                    Async::NotReady => {}
                                },
                                Async::NotReady => {}
                            }

                            self.state = RetryState::TimerActive(timer);
                            return Ok(AsyncSink::NotReady(item));
                        }
                        RetryPolicy::ForwardError(err) => return Err(err),
                    }
                }
                Ok(ok) => return Ok(ok),
            }
        }
    }
}

/// An extention trait for `Sink` which allows to use `SinkRetry` in a chain-like manner.
///
/// # Example
///
/// This magic trait allows you to handle errors on sink in a very neat manner:
/// ```
/// # use futures_retry::{RetryPolicy, SinkRetryExt, StreamRetryExt};
/// # use std::time::Duration;
/// # use tokio::io;
/// # use tokio::net::TcpListener;
/// # use tokio::prelude::*;
///
/// fn main() {
///      let addr = "127.0.0.1:12345".parse().unwrap();
///      let tcp = TcpListener::bind(&addr).unwrap();
///
///     let conn_error_handler = |e: io::Error| match e.kind() {
///         io::ErrorKind::Interrupted
///         | io::ErrorKind::ConnectionRefused
///         | io::ErrorKind::ConnectionReset
///         | io::ErrorKind::ConnectionAborted
///         | io::ErrorKind::NotConnected
///         | io::ErrorKind::BrokenPipe => RetryPolicy::Repeat,
///         io::ErrorKind::PermissionDenied => RetryPolicy::ForwardError(e),
///         _ => RetryPolicy::WaitRetry(Duration::from_millis(5)),
///     };
///
///     let data_sending_error_handler = |e: io::Error| match e.kind() {
///         io::ErrorKind::Interrupted => RetryPolicy::Repeat,
///         io::ErrorKind::TimedOut | io::ErrorKind::InvalidInput | io::ErrorKind::InvalidData => {
///             RetryPolicy::WaitRetry(Duration::from_millis(5))
///         }
///         _ => RetryPolicy::ForwardError(e),
///     };
///
///     let server = tcp
///         .incoming()
///         .retry(conn_error_handler)
///         .for_each(move |tcp| {
///             let (reader, writer) = tcp.split();
///
///             let reader = tokio::codec::FramedRead::new(reader, tokio::codec::LinesCodec::new());
///             let writer = tokio::codec::FramedWrite::new(writer, tokio::codec::LinesCodec::new());
///             // Copy the data back to the client
///
///             let conn = writer
///                 .retry(data_sending_error_handler) // retry
///                 .send_all(reader.retry(data_sending_error_handler))
///                 // when future is resolved sink and stream is returned we just drop them.
///                 .map(drop)
///                 // Handle any errors
///                 .map_err(|err| eprintln!("Can't copy data: IO error {:?}", err));
///
///             // Spawn the future as a concurrent task
///             tokio::spawn(conn);
///             Ok(())
///         })
///         .map_err(|err| {
///             eprintln!("server error {:?}", err);
///         });
///     tokio::run(server.select(Ok(())).map(|(_, _)| ()).map_err(|(_, _)| ()));
/// }
/// ```
pub trait SinkRetryExt: Sink {
    /// Converts the sink into a **retry sink**. See `SinkRetry::new` for details.
    ///
    /// # Warning (Implementation details)
    ///
    /// It depends from inner sink what happen when inner `start_send()` is resolve to error. This
    /// function assume that error has the same meaning as `AsyncSink::NotReady`. If your item will
    /// be buffered before error was returned - SinkRetry will insert this item again anyway.
    fn retry<F>(self, error_action: F) -> SinkRetry<F, Self>
    where
        Self: Sized,
        F: ErrorHandler<Self::SinkError>,
        Self::SinkItem: Clone,
        F: ErrorHandler<Self::SinkError>,
        F::OutError: From<Self::SinkError>,
    {
        SinkRetry::new(self, error_action)
    }
}

impl<S: ?Sized> SinkRetryExt for S where S: Sink {}

enum RetryState {
    WaitingForSink,
    TimerActive(tokio_timer::Delay),
}
impl RetryState {
    #[inline]
    fn is_waiting_for_sink(&self) -> bool {
        match self {
            RetryState::WaitingForSink => true,
            _ => false,
        }
    }
}

impl<F, S> Sink for SinkRetry<F, S>
where
    S: Sink,
    S::SinkItem: Clone,
    F: ErrorHandler<S::SinkError>,
    F::OutError: From<S::SinkError>,
{
    /// The type of value that the sink accepts.
    type SinkItem = S::SinkItem;

    /// The type of value produced by the sink when an error occurs.
    type SinkError = F::OutError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        match self.state {
            RetryState::WaitingForSink => self.try_send_item(item),
            RetryState::TimerActive(ref mut timer) => match timer.poll().expect("Timer panic!") {
                Async::NotReady => Ok(AsyncSink::NotReady(item)),
                Async::Ready(()) => self.try_send_item(item),
            },
        }
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        loop {
            let new_state = match self.state {
                //FIXME BROKEN RULE: DON'T REAPEAT YOURSELF -- the same code is in stream.rs!
                RetryState::TimerActive(ref mut delay) => match delay.poll() {
                    Ok(Async::Ready(())) => RetryState::WaitingForSink,
                    Ok(Async::NotReady) => return Ok(Async::NotReady),
                    Err(e) => {
                        // There could be two possible errors: timeout (TimerError::TooLong) or no
                        // new timer could be created (TimerError::NoCapacity).
                        // Since we are using the `sleep` method there could be no **timeout**
                        // error emitted.
                        // If the timer has reached its capacity.. well.. we are using just one
                        // timer.. so it will make me panic for sure.
                        panic!("Timer error: {}", e)
                    }
                },
                RetryState::WaitingForSink => match self.sink.poll_complete() {
                    Ok(x) => {
                        self.error_action.ok();
                        return Ok(x);
                    }
                    Err(e) => match self.error_action.handle(e) {
                        RetryPolicy::ForwardError(e) => return Err(e),
                        RetryPolicy::Repeat => RetryState::WaitingForSink,
                        RetryPolicy::WaitRetry(duration) => RetryState::TimerActive(
                            tokio_timer::Delay::new(Instant::now() + duration),
                        ),
                    },
                },
            };
            self.state = new_state;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::marker::PhantomData;

    enum SinkReturn {
        ReadyToFlush,
        NotReadyToFlush,
    }

    struct SinkIterResultMock<T, I> {
        iter: I,
        _sink_item: PhantomData<T>,
    }

    fn iter_result<T, J, E>(i: J) -> SinkIterResultMock<T, J::IntoIter>
    where
        J: IntoIterator<Item = Result<SinkReturn, E>>,
    {
        SinkIterResultMock {
            iter: i.into_iter(),
            _sink_item: PhantomData,
        }
    }

    impl<T, I, E> Sink for SinkIterResultMock<T, I>
    where
        I: Iterator<Item = Result<SinkReturn, E>>,
    {
        type SinkItem = T;
        type SinkError = E;

        fn start_send(
            &mut self,
            item: Self::SinkItem,
        ) -> StartSend<Self::SinkItem, Self::SinkError> {
            match self.poll_complete()? {
                Async::Ready(()) => Ok(AsyncSink::Ready),
                Async::NotReady => Ok(AsyncSink::NotReady(item)),
            }
        }

        fn poll_complete(&mut self) -> Poll<(), E> {
            match self.iter.next().expect("Iterator called after done!")? {
                SinkReturn::ReadyToFlush => Ok(Async::Ready(())),
                SinkReturn::NotReadyToFlush => Ok(Async::NotReady),
            }
        }
    }

    #[test]
    fn get_item_when_error_is_handling() {
        let sink = iter_result(vec![
            Ok(SinkReturn::NotReadyToFlush),
            Err(17u8),
            Ok(SinkReturn::NotReadyToFlush),
            Ok(SinkReturn::NotReadyToFlush),
        ]);
        let mut retry = sink.retry(|_| RetryPolicy::Repeat::<u8>);

        assert_eq!(Ok(AsyncSink::NotReady(5)), retry.start_send(5));
        assert_eq!(Ok(AsyncSink::NotReady(7)), retry.start_send(7));
        assert_eq!(Ok(AsyncSink::NotReady(8)), retry.start_send(8));
    }

    #[test]
    fn repeat() {
        let sink = iter_result::<u8, _, _>(vec![
            Ok(SinkReturn::ReadyToFlush),
            Err(17u64),
            Ok(SinkReturn::ReadyToFlush),
            Ok(SinkReturn::NotReadyToFlush),
        ]);
        let mut retry = SinkRetry::new(sink, |_| RetryPolicy::Repeat::<u64>);

        assert_eq!(Ok(AsyncSink::Ready), retry.start_send(2));
        assert_eq!(Ok(AsyncSink::Ready), retry.start_send(2));
        assert_eq!(Ok(AsyncSink::NotReady(2)), retry.start_send(2));
    }

    #[test]
    fn propagate() {
        let sink = iter_result::<u8, _, _>(vec![Err(17u8), Ok(SinkReturn::ReadyToFlush)]);
        let mut retry = SinkRetry::new(sink, RetryPolicy::ForwardError);
        assert_eq!(Err(17u8), retry.start_send(3));
        assert_eq!(Ok(AsyncSink::Ready), retry.start_send(3));
    }
}