tokio-w-custom-runtime.rs

// This example shows how to use the tokio runtime with any other executor
//
// The main components are a spawn fn that will wrap futures in a special future
// that will always enter the tokio context on poll. This only spawns on extra thread
// to manage and run the tokio drivers in the background.
//
// Playground link: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=22c3645ba6a03ec77459fd1dfc9f35f1

use tokio::net::TcpListener;
use tokio::sync::oneshot;

fn main() {
    let (tx, rx) = oneshot::channel();

    my_custom_runtime::spawn(async move {
        let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();

        println!("addr: {:?}", listener.local_addr());

        tx.send(()).unwrap();
    });

    futures::executor::block_on(rx).unwrap();
}

mod my_custom_runtime {
    use once_cell::sync::Lazy;
    use pin_project::pin_project;
    use std::{
        future::Future,
        pin::Pin,
        task::{Context, Poll},
    };
    use tokio::runtime::Handle;

    pub fn spawn(f: impl Future<Output = ()> + Send + 'static) {
        EXECUTOR.spawn(f);
    }

    struct ThreadPool {
        inner: futures::executor::ThreadPool,
        rt: tokio::runtime::Runtime,
    }

    static EXECUTOR: Lazy<ThreadPool> = Lazy::new(|| {
        // Spawn tokio runtime on a single background thread
        // enabling IO and timers.
        let rt = tokio::runtime::Builder::new()
            .basic_scheduler()
            .enable_all()
            .core_threads(1)
            .build()
            .unwrap();

        let inner = futures::executor::ThreadPool::builder().create().unwrap();

        ThreadPool { inner, rt }
    });

    impl ThreadPool {
        fn spawn(&self, f: impl Future<Output = ()> + Send + 'static) {
            let handle = self.rt.handle().clone();
            self.inner.spawn_ok(TokioIo { handle, inner: f });
        }
    }

    #[pin_project]
    struct TokioIo<F> {
        #[pin]
        inner: F,
        handle: Handle,
    }

    impl<F: Future> Future for TokioIo<F> {
        type Output = F::Output;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            let me = self.project();

            // Compiler is having trouble figuring out these
            // lifetimes.
            let handle = me.handle;
            let fut = me.inner;

            handle.enter(|| fut.poll(cx))
        }
    }
}