1
- use std:: { future:: Future , thread } ;
1
+ use std:: { future:: Future , sync :: Mutex } ;
2
2
3
3
use :: tokio:: {
4
4
runtime:: { Builder , Runtime } ,
5
5
task,
6
6
} ;
7
- use futures:: future:: pending;
8
- use once_cell:: sync:: OnceCell ;
7
+ use once_cell:: sync:: { Lazy , OnceCell } ;
9
8
use pyo3:: prelude:: * ;
10
9
11
10
use crate :: generic;
@@ -30,10 +29,9 @@ pub use pyo3_asyncio_macros::tokio_main as main;
30
29
#[ cfg( all( feature = "attributes" , feature = "testing" ) ) ]
31
30
pub use pyo3_asyncio_macros:: tokio_test as test;
32
31
32
+ static TOKIO_BUILDER : Lazy < Mutex < Builder > > = Lazy :: new ( || Mutex :: new ( multi_thread ( ) ) ) ;
33
33
static TOKIO_RUNTIME : OnceCell < Runtime > = OnceCell :: new ( ) ;
34
34
35
- const EXPECT_TOKIO_INIT : & str = "Tokio runtime must be initialized" ;
36
-
37
35
impl generic:: JoinError for task:: JoinError {
38
36
fn is_panic ( & self ) -> bool {
39
37
task:: JoinError :: is_panic ( self )
@@ -65,79 +63,26 @@ impl generic::SpawnLocalExt for TokioRuntime {
65
63
}
66
64
}
67
65
68
- /// Initialize the Tokio Runtime with a custom build
69
- pub fn init ( runtime : Runtime ) {
70
- TOKIO_RUNTIME
71
- . set ( runtime)
72
- . expect ( "Tokio Runtime has already been initialized" ) ;
73
- }
74
-
75
- fn current_thread ( ) -> Runtime {
76
- Builder :: new_current_thread ( )
77
- . enable_all ( )
78
- . build ( )
79
- . expect ( "Couldn't build the current-thread Tokio runtime" )
80
- }
81
-
82
- fn start_current_thread ( ) {
83
- thread:: spawn ( move || {
84
- TOKIO_RUNTIME . get ( ) . unwrap ( ) . block_on ( pending :: < ( ) > ( ) ) ;
85
- } ) ;
86
- }
87
-
88
- /// Initialize the Tokio Runtime with current-thread scheduler
89
- ///
90
- /// # Panics
91
- /// This function will panic if called a second time. See [`init_current_thread_once`] if you want
92
- /// to avoid this panic.
93
- pub fn init_current_thread ( ) {
94
- init ( current_thread ( ) ) ;
95
- start_current_thread ( ) ;
66
+ /// Initialize the Tokio runtime with a custom build
67
+ pub fn init ( builder : Builder ) {
68
+ * TOKIO_BUILDER . lock ( ) . unwrap ( ) = builder
96
69
}
97
70
98
71
/// Get a reference to the current tokio runtime
99
72
pub fn get_runtime < ' a > ( ) -> & ' a Runtime {
100
- TOKIO_RUNTIME . get ( ) . expect ( EXPECT_TOKIO_INIT )
101
- }
102
-
103
- fn multi_thread ( ) -> Runtime {
104
- Builder :: new_multi_thread ( )
105
- . enable_all ( )
106
- . build ( )
107
- . expect ( "Couldn't build the multi-thread Tokio runtime" )
108
- }
109
-
110
- /// Initialize the Tokio Runtime with the multi-thread scheduler
111
- ///
112
- /// # Panics
113
- /// This function will panic if called a second time. See [`init_multi_thread_once`] if you want to
114
- /// avoid this panic.
115
- pub fn init_multi_thread ( ) {
116
- init ( multi_thread ( ) ) ;
117
- }
118
-
119
- /// Ensure that the Tokio Runtime is initialized
120
- ///
121
- /// If the runtime has not been initialized already, the multi-thread scheduler
122
- /// is used. Calling this function a second time is a no-op.
123
- pub fn init_multi_thread_once ( ) {
124
- TOKIO_RUNTIME . get_or_init ( || multi_thread ( ) ) ;
125
- }
126
-
127
- /// Ensure that the Tokio Runtime is initialized
128
- ///
129
- /// If the runtime has not been initialized already, the current-thread
130
- /// scheduler is used. Calling this function a second time is a no-op.
131
- pub fn init_current_thread_once ( ) {
132
- let mut initialized = false ;
133
73
TOKIO_RUNTIME . get_or_init ( || {
134
- initialized = true ;
135
- current_thread ( )
136
- } ) ;
74
+ TOKIO_BUILDER
75
+ . lock ( )
76
+ . unwrap ( )
77
+ . build ( )
78
+ . expect ( "Unable to build Tokio runtime" )
79
+ } )
80
+ }
137
81
138
- if initialized {
139
- start_current_thread ( ) ;
140
- }
82
+ fn multi_thread ( ) -> Builder {
83
+ let mut builder = Builder :: new_multi_thread ( ) ;
84
+ builder. enable_all ( ) ;
85
+ builder
141
86
}
142
87
143
88
/// Run the event loop until the given Future completes
@@ -157,16 +102,9 @@ pub fn init_current_thread_once() {
157
102
/// # use std::time::Duration;
158
103
/// #
159
104
/// # use pyo3::prelude::*;
160
- /// # use tokio::runtime::{Builder, Runtime};
161
- /// #
162
- /// # let runtime = Builder::new_current_thread()
163
- /// # .enable_all()
164
- /// # .build()
165
- /// # .expect("Couldn't build the runtime");
166
105
/// #
167
106
/// # Python::with_gil(|py| {
168
107
/// # pyo3_asyncio::with_runtime(py, || {
169
- /// # pyo3_asyncio::tokio::init_current_thread();
170
108
/// pyo3_asyncio::tokio::run_until_complete(py, async move {
171
109
/// tokio::time::sleep(Duration::from_secs(1)).await;
172
110
/// Ok(())
0 commit comments