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

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>,
}

/// A guard that resets the current timer to `None` when dropped.
#[derive(Debug)]
pub struct DefaultGuard {
    _p: (),
}

thread_local! {
    /// Tracks the timer for the current execution context.
    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,
{
    let _guard = set_default(handle);
    f(enter)
}

/// Sets `handle` as the default timer, returning a guard that unsets it on drop.
///
/// # Panics
///
/// This function panics if there already is a default timer set.
pub fn set_default(handle: &Handle) -> DefaultGuard {
    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());
    });
    DefaultGuard { _p: () }
}

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),
        }
    }

    #[doc(hidden)]
    #[deprecated(since = "0.2.11", note = "use timeout instead")]
    pub fn deadline<T>(&self, future: T, deadline: Instant) -> Deadline<T> {
        Deadline::new_with_delay(future, self.delay(deadline))
    }

    /// Create a `Timeout` driven by this handle's associated `Timer`.
    pub fn timeout<T>(&self, value: T, deadline: Instant) -> Timeout<T> {
        Timeout::new_with_delay(value, 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")
    }
}

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