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
use {Error, Delay, Deadline, Interval};
use timer::Inner;

use tokio_executor::Enter;

use std::cell::RefCell;
use std::fmt;
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};

/// Handle to timer instance.
///
/// The `Handle` allows creating `Delay` instances that are driven by the
/// associated timer.
///
/// A `Handle` is obtained by calling [`Timer::handle`], [`Handle::current`], or
/// [`Handle::default`].
///
/// * [`Timer::handle`]: returns a handle associated with the specific timer.
///   The handle will always reference the same timer.
///
/// * [`Handle::current`]: returns a handle to the timer for the execution
///   context **at the time the function is called**. This function must be
///   called from a runtime that has an associated timer or it will panic.
///   The handle will always reference the same timer.
///
/// * [`Handle::default`]: returns a handle to the timer for the execution
///   context **at the time the handle is used**. This function is safe to call
///   at any time. The handle may reference different specific timer instances.
///   Calling `Handle::default().delay(...)` is always equivalent to
///   `Delay::new(...)`.
///
/// [`Timer::handle`]: struct.Timer.html#method.handle
/// [`Handle::current`]: #method.current
/// [`Handle::default`]: #method.default
#[derive(Debug, Clone)]
pub struct Handle {
    inner: Option<HandlePriv>,
}

/// Like `Handle` but never `None`.
#[derive(Clone)]
pub(crate) struct HandlePriv {
    inner: Weak<Inner>,
}

/// Tracks the timer for the current execution context.
thread_local!(static CURRENT_TIMER: RefCell<Option<HandlePriv>> = RefCell::new(None));

/// Set the default timer for the duration of the closure.
///
/// From within the closure, [`Delay`] instances that are created via
/// [`Delay::new`] can be used.
///
/// # Panics
///
/// This function panics if there already is a default timer set.
///
/// [`Delay`]: ../struct.Delay.html
/// [`Delay::new`]: ../struct.Delay.html#method.new
pub fn with_default<F, R>(handle: &Handle, enter: &mut Enter, f: F) -> R
where F: FnOnce(&mut Enter) -> R
{
    // Ensure that the timer is removed from the thread-local context
    // when leaving the scope. This handles cases that involve panicking.
    struct Reset;

    impl Drop for Reset {
        fn drop(&mut self) {
            CURRENT_TIMER.with(|current| {
                let mut current = current.borrow_mut();
                *current = None;
            });
        }
    }

    // This ensures the value for the current timer gets reset even if there is
    // a panic.
    let _r = Reset;

    CURRENT_TIMER.with(|current| {
        {
            let mut current = current.borrow_mut();

            assert!(current.is_none(), "default Tokio timer already set \
                    for execution context");

            let handle = handle.as_priv()
                .unwrap_or_else(|| panic!("`handle` does not reference a timer"));

            *current = Some(handle.clone());
        }

        f(enter)
    })
}

impl Handle {
    pub(crate) fn new(inner: Weak<Inner>) -> Handle {
        let inner = HandlePriv { inner };
        Handle { inner: Some(inner) }
    }

    /// Returns a handle to the current timer.
    ///
    /// The current timer is the timer that is currently set as default using
    /// [`with_default`].
    ///
    /// This function should only be called from within the context of
    /// [`with_default`]. Calling this function from outside of this context
    /// will return a `Handle` that does not reference a timer. `Delay`
    /// instances created with this handle will error.
    ///
    /// See [type] level documentation for more ways to obtain a `Handle` value.
    ///
    /// [`with_default`]: ../fn.with_default.html
    /// [type]: #
    pub fn current() -> Handle {
        let private = HandlePriv::try_current()
            .unwrap_or_else(|_| {
                HandlePriv { inner: Weak::new() }
            });

        Handle { inner: Some(private) }
    }

    /// Create a `Delay` driven by this handle's associated `Timer`.
    pub fn delay(&self, deadline: Instant) -> Delay {
        match self.inner {
            Some(ref handle_priv) => {
                Delay::new_with_handle(deadline, handle_priv.clone())
            }
            None => {
                Delay::new(deadline)
            }
        }
    }

    /// Create a `Deadline` driven by this handle's associated `Timer`.
    pub fn deadline<T>(&self, future: T, deadline: Instant) -> Deadline<T> {
        Deadline::new_with_delay(future, self.delay(deadline))
    }

    /// Create a new `Interval` that starts at `at` and yields every `duration`
    /// interval after that.
    pub fn interval(&self, at: Instant, duration: Duration) -> Interval {
        Interval::new_with_delay(self.delay(at), duration)
    }

    fn as_priv(&self) -> Option<&HandlePriv> {
        self.inner.as_ref()
    }
}

impl Default for Handle {
    fn default() -> Handle {
        Handle { inner: None }
    }
}

impl HandlePriv {
    /// Try to get a handle to the current timer.
    ///
    /// Returns `Err` if no handle is found.
    pub(crate) fn try_current() -> Result<HandlePriv, Error> {
        CURRENT_TIMER.with(|current| {
            match *current.borrow() {
                Some(ref handle) => Ok(handle.clone()),
                None => Err(Error::shutdown()),
            }
        })
    }

    /// Try to return a strong ref to the inner
    pub(crate) fn inner(&self) -> Option<Arc<Inner>> {
        self.inner.upgrade()
    }

    /// Consume the handle, returning the weak Inner ref.
    pub(crate) fn into_inner(self) -> Weak<Inner> {
        self.inner
    }
}

impl fmt::Debug for HandlePriv {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HandlePriv")
    }
}