Merge pull request #14 from awestlake87/init-changes

Init changes for #13

Author: awestlake87

Cargo.toml

@@ -98,6 +98,6 @@ features = ["unstable"]
 optional = true
 
 [dependencies.tokio]
-version = "1.0" 
+version = "1.4" 
 features = ["full"]
 optional = true

pytests/common/mod.rs

@@ -51,3 +51,11 @@ pub(super) async fn test_other_awaitables() -> PyResult<()> {
 
     Ok(())
 }
+
+pub(super) fn test_init_twice() -> PyResult<()> {
+    // try_init has already been called in test main - ensure a second call doesn't mess the other
+    // tests up
+    Python::with_gil(|py| pyo3_asyncio::try_init(py))?;
+
+    Ok(())
+}

pytests/test_async_std_asyncio.rs

@@ -69,6 +69,11 @@ async fn test_other_awaitables() -> PyResult<()> {
     common::test_other_awaitables().await
 }
 
+#[pyo3_asyncio::async_std::test]
+fn test_init_twice() -> PyResult<()> {
+    common::test_init_twice()
+}
+
 #[pyo3_asyncio::async_std::main]
 async fn main() -> pyo3::PyResult<()> {
     pyo3_asyncio::testing::main().await

pytests/tokio_asyncio/mod.rs

@@ -67,3 +67,18 @@ async fn test_into_future() -> PyResult<()> {
 async fn test_other_awaitables() -> PyResult<()> {
     common::test_other_awaitables().await
 }
+
+#[pyo3_asyncio::tokio::test]
+fn test_init_twice() -> PyResult<()> {
+    common::test_init_twice()
+}
+
+#[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(())
+}

src/lib.rs

@@ -193,29 +193,29 @@ where
 
 /// Attempt to initialize the Python and Rust event loops
 ///
-/// Must be called at the start of your program
-fn try_init(py: Python) -> PyResult<()> {
-    let asyncio = py.import("asyncio")?;
-
-    let ensure_future = asyncio.getattr("ensure_future")?;
-
-    let event_loop = asyncio.call_method0("get_event_loop")?;
-    let executor = py
-        .import("concurrent.futures.thread")?
-        .getattr("ThreadPoolExecutor")?
-        .call0()?;
-
-    event_loop.call_method1("set_default_executor", (executor,))?;
-
-    let call_soon = event_loop.getattr("call_soon_threadsafe")?;
-    let create_future = event_loop.getattr("create_future")?;
-
-    ASYNCIO.get_or_init(|| asyncio.into());
-    ENSURE_FUTURE.get_or_init(|| ensure_future.into());
-    EVENT_LOOP.get_or_init(|| event_loop.into());
-    EXECUTOR.get_or_init(|| executor.into());
-    CALL_SOON.get_or_init(|| call_soon.into());
-    CREATE_FUTURE.get_or_init(|| create_future.into());
+/// - Must be called before any other pyo3-asyncio functions.
+/// - Calling `try_init` a second time returns `Ok(())` and does nothing.
+///   > In future versions this may return an `Err`.
+pub fn try_init(py: Python) -> PyResult<()> {
+    EVENT_LOOP.get_or_try_init(|| -> PyResult<PyObject> {
+        let asyncio = py.import("asyncio")?;
+        let ensure_future = asyncio.getattr("ensure_future")?;
+        let event_loop = asyncio.call_method0("get_event_loop")?;
+        let executor = py
+            .import("concurrent.futures.thread")?
+            .getattr("ThreadPoolExecutor")?
+            .call0()?;
+        event_loop.call_method1("set_default_executor", (executor,))?;
+        let call_soon = event_loop.getattr("call_soon_threadsafe")?;
+        let create_future = event_loop.getattr("create_future")?;
+
+        ASYNCIO.get_or_init(|| asyncio.into());
+        ENSURE_FUTURE.get_or_init(|| ensure_future.into());
+        EXECUTOR.get_or_init(|| executor.into());
+        CALL_SOON.get_or_init(|| call_soon.into());
+        CREATE_FUTURE.get_or_init(|| create_future.into());
+        Ok(event_loop.into())
+    })?;
 
     Ok(())
 }
@@ -284,7 +284,7 @@ pub fn run_forever(py: Python) -> PyResult<()> {
 }
 
 /// Shutdown the event loops and perform any necessary cleanup
-fn try_close(py: Python) -> PyResult<()> {
+pub fn try_close(py: Python) -> PyResult<()> {
     // Shutdown the executor and wait until all threads are cleaned up
     EXECUTOR
         .get()

src/tokio.rs

@@ -63,33 +63,72 @@ pub fn init(runtime: 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(
-        Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .expect("Couldn't build the current-thread Tokio runtime"),
-    );
-
-    thread::spawn(|| {
-        get_runtime().block_on(pending::<()>());
-    });
+    init(current_thread());
+    start_current_thread();
 }
 
 /// 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(
-        Builder::new_multi_thread()
-            .enable_all()
-            .build()
-            .expect("Couldn't build the multi-thread Tokio runtime"),
-    );
+    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()
+    });
+
+    if initialized {
+        start_current_thread();
+    }
 }
 
 /// Run the event loop until the given Future completes