Merged in lazy tokio initialization

Author: Andrew J Westlake

README.md

@@ -124,8 +124,8 @@ Export an async function that makes use of `async-std`:
 use pyo3::{prelude::*, wrap_pyfunction};
 
 #[pyfunction]
-fn rust_sleep(py: Python) -> PyResult<PyObject> {
-    pyo3_asyncio::async_std::into_coroutine(py, async {
+fn rust_sleep(py: Python) -> PyResult<&PyAny> {
+    pyo3_asyncio::async_std::future_into_py(py, async {
         async_std::task::sleep(std::time::Duration::from_secs(1)).await;
         Ok(Python::with_gil(|py| py.None()))
     })
@@ -148,19 +148,15 @@ If you want to use `tokio` instead, here's what your module should look like:
 use pyo3::{prelude::*, wrap_pyfunction};
 
 #[pyfunction]
-fn rust_sleep(py: Python) -> PyResult<PyObject> {
-    pyo3_asyncio::tokio::into_coroutine(py, async {
+fn rust_sleep(py: Python) -> PyResult<&PyAny> {
+    pyo3_asyncio::tokio::future_into_py(py, async {
         tokio::time::sleep(std::time::Duration::from_secs(1)).await;
         Ok(Python::with_gil(|py| py.None()))
     })
 }
 
 #[pymodule]
 fn my_async_module(py: Python, m: &PyModule) -> PyResult<()> {
-    // Tokio needs explicit initialization before any pyo3-asyncio conversions.
-    // The module import is a prime place to do this.
-    pyo3_asyncio::tokio::init_multi_thread_once();
-
     m.add_function(wrap_pyfunction!(rust_sleep, m)?)?;
 
     Ok(())

pyo3-asyncio-macros/src/tokio.rs

@@ -219,16 +219,23 @@ fn parse_knobs(
 
     let config = config.build()?;
 
-    let mut rt = match config.flavor {
+    let builder = match config.flavor {
         RuntimeFlavor::CurrentThread => quote! {
             pyo3_asyncio::tokio::re_exports::runtime::Builder::new_current_thread()
         },
         RuntimeFlavor::Threaded => quote! {
             pyo3_asyncio::tokio::re_exports::runtime::Builder::new_multi_thread()
         },
     };
+
+    let mut builder_init = quote! {
+        builder.enable_all();
+    };
     if let Some(v) = config.worker_threads {
-        rt = quote! { #rt.worker_threads(#v) };
+        builder_init = quote! {
+            builder.worker_threads(#v);
+            #builder_init;
+        };
     }
 
     let rt_init = match config.flavor {
@@ -249,12 +256,10 @@ fn parse_knobs(
 
             pyo3::prepare_freethreaded_python();
 
-            pyo3_asyncio::tokio::init(
-                #rt
-                    .enable_all()
-                    .build()
-                    .unwrap()
-            );
+            let mut builder = #builder;
+            #builder_init;
+
+            pyo3_asyncio::tokio::init(builder);
 
             #rt_init
 

pytests/test_tokio_current_thread_asyncio.rs

@@ -10,7 +10,15 @@ fn main() -> pyo3::PyResult<()> {
     Python::with_gil(|py| {
         // into_coroutine requires the 0.13 API
         pyo3_asyncio::try_init(py)?;
-        pyo3_asyncio::tokio::init_current_thread();
+
+        let mut builder = tokio::runtime::Builder::new_current_thread();
+        builder.enable_all();
+
+        pyo3_asyncio::tokio::init(builder);
+        std::thread::spawn(move || {
+            pyo3_asyncio::tokio::get_runtime().block_on(futures::future::pending::<()>());
+        });
+
         pyo3_asyncio::tokio::run(py, pyo3_asyncio::testing::main())
     })
 }

pytests/test_tokio_current_thread_run_forever.rs

@@ -2,7 +2,14 @@ mod tokio_run_forever;
 
 fn main() {
     pyo3::prepare_freethreaded_python();
-    pyo3_asyncio::tokio::init_current_thread();
+
+    let mut builder = tokio::runtime::Builder::new_current_thread();
+    builder.enable_all();
+
+    pyo3_asyncio::tokio::init(builder);
+    std::thread::spawn(move || {
+        pyo3_asyncio::tokio::get_runtime().block_on(futures::future::pending::<()>());
+    });
 
     tokio_run_forever::test_main();
 }

pytests/test_tokio_multi_thread_asyncio.rs

@@ -10,7 +10,6 @@ fn main() -> pyo3::PyResult<()> {
     Python::with_gil(|py| {
         // into_coroutine requires the 0.13 API
         pyo3_asyncio::try_init(py)?;
-        pyo3_asyncio::tokio::init_multi_thread();
         pyo3_asyncio::tokio::run(py, pyo3_asyncio::testing::main())
     })
 }

pytests/test_tokio_multi_thread_run_forever.rs

@@ -2,8 +2,5 @@ mod tokio_run_forever;
 
 fn main() {
     pyo3::prepare_freethreaded_python();
-
-    pyo3_asyncio::tokio::init_multi_thread();
-
     tokio_run_forever::test_main();
 }

pytests/tokio_asyncio/mod.rs

@@ -124,16 +124,6 @@ async fn test_other_awaitables() -> PyResult<()> {
     .await
 }
 
-#[pyo3_asyncio::tokio::test]
-fn test_init_tokio_twice() -> PyResult<()> {
-    // tokio has already been initialized in test main. call these functions to
-    // make sure they don't cause problems with the other tests.
-    pyo3_asyncio::tokio::init_multi_thread_once();
-    pyo3_asyncio::tokio::init_current_thread_once();
-
-    Ok(())
-}
-
 #[pyo3_asyncio::tokio::test]
 fn test_local_future_into_py(event_loop: PyObject) -> PyResult<()> {
     tokio::task::LocalSet::new().block_on(pyo3_asyncio::tokio::get_runtime(), async {

src/tokio.rs

@@ -1,11 +1,13 @@
-use std::{future::Future, pin::Pin, thread};
+use std::{future::Future, pin::Pin, sync::Mutex};
 
 use ::tokio::{
     runtime::{Builder, Runtime},
     task,
 };
-use futures::future::pending;
-use once_cell::{sync::OnceCell, unsync::OnceCell as UnsyncOnceCell};
+use once_cell::{
+    sync::{Lazy, OnceCell},
+    unsync::OnceCell as UnsyncOnceCell,
+};
 use pyo3::prelude::*;
 
 use crate::generic::{self, Runtime as GenericRuntime, SpawnLocalExt};
@@ -30,10 +32,9 @@ pub use pyo3_asyncio_macros::tokio_main as main;
 #[cfg(all(feature = "attributes", feature = "testing"))]
 pub use pyo3_asyncio_macros::tokio_test as test;
 
+static TOKIO_BUILDER: Lazy<Mutex<Builder>> = Lazy::new(|| Mutex::new(multi_thread()));
 static TOKIO_RUNTIME: OnceCell<Runtime> = OnceCell::new();
 
-const EXPECT_TOKIO_INIT: &str = "Tokio runtime must be initialized";
-
 impl generic::JoinError for task::JoinError {
     fn is_panic(&self) -> bool {
         task::JoinError::is_panic(self)
@@ -121,79 +122,26 @@ pub fn get_current_loop(py: Python) -> PyResult<&PyAny> {
     generic::get_current_loop::<TokioRuntime>(py)
 }
 
-/// Initialize the Tokio Runtime with a custom build
-pub fn init(runtime: Runtime) {
-    TOKIO_RUNTIME
-        .set(runtime)
-        .expect("Tokio Runtime has already been initialized");
-}
-
-fn current_thread() -> Runtime {
-    Builder::new_current_thread()
-        .enable_all()
-        .build()
-        .expect("Couldn't build the current-thread Tokio runtime")
-}
-
-fn start_current_thread() {
-    thread::spawn(move || {
-        TOKIO_RUNTIME.get().unwrap().block_on(pending::<()>());
-    });
-}
-
-/// Initialize the Tokio Runtime with current-thread scheduler
-///
-/// # Panics
-/// This function will panic if called a second time. See [`init_current_thread_once`] if you want
-/// to avoid this panic.
-pub fn init_current_thread() {
-    init(current_thread());
-    start_current_thread();
+/// Initialize the Tokio runtime with a custom build
+pub fn init(builder: Builder) {
+    *TOKIO_BUILDER.lock().unwrap() = builder
 }
 
 /// Get a reference to the current tokio runtime
 pub fn get_runtime<'a>() -> &'a Runtime {
-    TOKIO_RUNTIME.get().expect(EXPECT_TOKIO_INIT)
-}
-
-fn multi_thread() -> Runtime {
-    Builder::new_multi_thread()
-        .enable_all()
-        .build()
-        .expect("Couldn't build the multi-thread Tokio runtime")
-}
-
-/// Initialize the Tokio Runtime with the multi-thread scheduler
-///
-/// # Panics
-/// This function will panic if called a second time. See [`init_multi_thread_once`] if you want to
-/// avoid this panic.
-pub fn init_multi_thread() {
-    init(multi_thread());
-}
-
-/// Ensure that the Tokio Runtime is initialized
-///
-/// If the runtime has not been initialized already, the multi-thread scheduler
-/// is used. Calling this function a second time is a no-op.
-pub fn init_multi_thread_once() {
-    TOKIO_RUNTIME.get_or_init(|| multi_thread());
-}
-
-/// Ensure that the Tokio Runtime is initialized
-///
-/// If the runtime has not been initialized already, the current-thread
-/// scheduler is used. Calling this function a second time is a no-op.
-pub fn init_current_thread_once() {
-    let mut initialized = false;
     TOKIO_RUNTIME.get_or_init(|| {
-        initialized = true;
-        current_thread()
-    });
+        TOKIO_BUILDER
+            .lock()
+            .unwrap()
+            .build()
+            .expect("Unable to build Tokio runtime")
+    })
+}
 
-    if initialized {
-        start_current_thread();
-    }
+fn multi_thread() -> Builder {
+    let mut builder = Builder::new_multi_thread();
+    builder.enable_all();
+    builder
 }
 
 /// Run the event loop until the given Future completes
@@ -213,17 +161,10 @@ pub fn init_current_thread_once() {
 /// # use std::time::Duration;
 /// #
 /// # use pyo3::prelude::*;
-/// # use tokio::runtime::{Builder, Runtime};
-/// #
-/// # let runtime = Builder::new_current_thread()
-/// #     .enable_all()
-/// #     .build()
-/// #     .expect("Couldn't build the runtime");
 /// #
 /// # pyo3::prepare_freethreaded_python();
 /// # Python::with_gil(|py| {
 /// # pyo3_asyncio::with_runtime(py, || {
-/// # pyo3_asyncio::tokio::init_current_thread();
 /// # let event_loop = py.import("asyncio")?.call_method0("new_event_loop")?;
 /// pyo3_asyncio::tokio::run_until_complete(event_loop, async move {
 ///     tokio::time::sleep(Duration::from_secs(1)).await;