fix race condition

python/src/utils.rs

@@ -11,21 +11,16 @@ use tokio::runtime::Runtime;
 pub fn rt() -> &'static Runtime {
     static TOKIO_RT: OnceLock<Runtime> = OnceLock::new();
     static PID: OnceLock<u32> = 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."))
 }