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

//!Single threaded runtime module
//!
//!Runtime is stored in thread local variable
//!## Example
//!
//!```rust
//!use tokio_global::single as rt;
//!use tokio_global::AutoRuntime;
//!
//!let _guard = rt::init();
//!//Now we can exeute futures using runtime
//!//When guard goes out of scope though,
//!//we no longer can use it.
//!
//!let result = futures::future::ok::<u32, u32>(5).finish().expect("Ok");
//!assert_eq!(result, 5);
//!
//!```


use tokio::runtime::current_thread::{self, Handle};
use futures::{IntoFuture, Future};

use std::io;
use std::cell::Cell;
use std::marker::PhantomData;

const RUNTIME_NOT_AVAIL: &'static str = "Runtime is not available! Initialize it or do not use it within blocking calls!";

thread_local!(static TOKIO: Cell<Option<current_thread::Runtime>> = Cell::new(None));

///Guard that controls lifetime of Runtime module
///
///On drop, it terminates runtime making it impossible to use it any longer.
pub struct Runtime {
    //workaround for unstable !Send
    _dummy: PhantomData<*mut u8>
}

impl Runtime {
    ///Initializes new runtime and returns guard that controls its lifetime.
    ///
    ///This function must be called prior to any usage of runtime related functionality.
    ///
    ///## Panics
    ///
    ///If runtime is already initialized.
    pub fn new() -> io::Result<Self> {
        let runtime = current_thread::Runtime::new()?;

        TOKIO.with(|rt| match rt.replace(None) {
            None => rt.set(Some(runtime)),
            Some(old) => {
                drop(old);
                panic!("Double set of runtime!")
            }
        });

        Ok(Runtime {
            _dummy: PhantomData
        })
    }
}

impl Drop for Runtime {
    fn drop(&mut self) {
        TOKIO.with(|rt| rt.replace(None));
    }
}

///Initializes new runtime and returns guard that controls its lifetime.
///
///This function must be called prior to any usage of runtime related functionality.
///
///## Panics
///
///If runtime is already initialized.
pub fn init() -> Runtime {
    Runtime::new().expect("To create runtime module")
}

///Starts function within tokio runtime that returns future
///and awaits its completion.
///
///## Note
///
///It must not be used from within async context
pub fn run<R: IntoFuture, F: FnOnce() -> R>(runner: F) -> Result<R::Item, R::Error> {
    block_on(runner().into_future())
}

///Run a future to completion on the Tokio runtime.
///
///This runs the given future on the runtime, blocking until it is complete,
///and yielding its resolved result.
///Any tasks or timers which the future spawns internally will be executed on the runtime.
///
///## Note
///
///It must not be used from within async context
pub fn block_on<F: Future>(future: F) -> Result<F::Item, F::Error> {
    TOKIO.with(|rt| match unsafe { &mut *rt.as_ptr() } {
        Some(tokio) => tokio.block_on(future),
        None => panic!(RUNTIME_NOT_AVAIL)
    })

}

///Spawns future on runtime's event loop.
pub fn spawn<F: Future<Item=(), Error=()> + 'static>(fut: F) {
    TOKIO.with(|rt| match unsafe { &mut *rt.as_ptr() } {
        Some(tokio) => tokio.spawn(fut),
        None => panic!(RUNTIME_NOT_AVAIL)
    });
}

///Retrieves tokio's handle.
pub fn handle() -> Handle {
    TOKIO.with(|rt| match unsafe { &*rt.as_ptr() } {
        Some(tokio) => tokio.handle(),
        None => panic!(RUNTIME_NOT_AVAIL)
    })
}

#[cfg(not(feature = "multi"))]
impl<F: Future> super::AutoRuntime for F {
    #[inline]
    fn finish(self) -> Result<Self::Item, Self::Error> {
        block_on(self)
    }

    #[inline]
    fn spawn(self) where Self: 'static + Future<Item=(), Error=()> {
        spawn(self);
    }
}