From 5b1ab37adb614171bbc327728db99530977ff135 Mon Sep 17 00:00:00 2001 From: Peter Ke Date: Thu, 10 Oct 2024 19:57:27 -0700 Subject: [PATCH] fix race condition --- python/src/utils.rs | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/python/src/utils.rs b/python/src/utils.rs index b063b64d08..0e31682b54 100644 --- a/python/src/utils.rs +++ b/python/src/utils.rs @@ -11,21 +11,16 @@ use tokio::runtime::Runtime; pub fn rt() -> &'static Runtime { static TOKIO_RT: OnceLock = OnceLock::new(); static PID: OnceLock = OnceLock::new(); - match PID.get() { - Some(pid) if pid == &std::process::id() => {} // Reuse the static runtime. - Some(pid) => { - panic!( - "Forked process detected - current PID is {} but the tokio runtime was created by {}. The tokio \ - runtime does not support forked processes https://github.com/tokio-rs/tokio/issues/4301. If you are \ - seeing this message while using Python multithreading make sure to use the `spawn` or `forkserver` \ - mode.", - pid, std::process::id() - ); - } - None => { - PID.set(std::process::id()) - .expect("Failed to record PID for tokio runtime."); - } + let pid = std::process::id(); + let runtime_pid = *PID.get_or_init(|| pid); + if pid != runtime_pid { + panic!( + "Forked process detected - current PID is {} but the tokio runtime was created by {}. The tokio \ + runtime does not support forked processes https://github.com/tokio-rs/tokio/issues/4301. If you are \ + seeing this message while using Python multithreading make sure to use the `spawn` or `forkserver` \ + mode.", + pid, runtime_pid + ); } TOKIO_RT.get_or_init(|| Runtime::new().expect("Failed to create a tokio runtime.")) }