From 0b8ec14022451789d734a1d32c80fe16cd36dbbf Mon Sep 17 00:00:00 2001
From: Ion Koutsouris <15728914+ion-elgreco@users.noreply.github.com>
Date: Sat, 31 May 2025 17:36:26 +0200
Subject: [PATCH] feat: prevent fork with tokio

Signed-off-by: Ion Koutsouris <15728914+ion-elgreco@users.noreply.github.com>
---
 Cargo.lock                     | 312 ++++++++++++++++++++++++++++++++-
 Cargo.toml                     |   2 +-
 bindings/python/src/runtime.rs |  25 ++-
 3 files changed, 328 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6528d13f6..de15c83b3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2,6 +2,54 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "abi_stable"
+version = "0.11.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69d6512d3eb05ffe5004c59c206de7f99c34951504056ce23fc953842f12c445"
+dependencies = [
+ "abi_stable_derive",
+ "abi_stable_shared",
+ "const_panic",
+ "core_extensions",
+ "crossbeam-channel",
+ "generational-arena",
+ "libloading 0.7.4",
+ "lock_api",
+ "parking_lot",
+ "paste",
+ "repr_offset",
+ "rustc_version",
+ "serde",
+ "serde_derive",
+ "serde_json",
+]
+
+[[package]]
+name = "abi_stable_derive"
+version = "0.11.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d7178468b407a4ee10e881bc7a328a65e739f0863615cca4429d43916b05e898"
+dependencies = [
+ "abi_stable_shared",
+ "as_derive_utils",
+ "core_extensions",
+ "proc-macro2",
+ "quote",
+ "rustc_version",
+ "syn 1.0.109",
+ "typed-arena",
+]
+
+[[package]]
+name = "abi_stable_shared"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2b5df7688c123e63f4d4d649cba63f2967ba7f7861b1664fca3f77d3dad2b63"
+dependencies = [
+ "core_extensions",
+]
+
 [[package]]
 name = "addr2line"
 version = "0.24.2"
@@ -227,6 +275,7 @@ dependencies = [
  "arrow-schema",
  "arrow-select",
  "arrow-string",
+ "pyo3",
 ]
 
 [[package]]
@@ -387,6 +436,9 @@ name = "arrow-schema"
 version = "55.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "73a47aa0c771b5381de2b7f16998d351a6f4eb839f1e13d48353e17e873d969b"
+dependencies = [
+ "bitflags",
+]
 
 [[package]]
 name = "arrow-select"
@@ -419,6 +471,18 @@ dependencies = [
  "regex-syntax 0.8.5",
 ]
 
+[[package]]
+name = "as_derive_utils"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff3c96645900a44cf11941c111bd08a6573b0e2f9f69bc9264b179d8fae753c4"
+dependencies = [
+ "core_extensions",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "assert-json-diff"
 version = "2.0.2"
@@ -504,6 +568,15 @@ dependencies = [
  "slab",
 ]
 
+[[package]]
+name = "async-ffi"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50"
+dependencies = [
+ "abi_stable",
+]
+
 [[package]]
 name = "async-global-executor"
 version = "2.4.1"
@@ -1457,7 +1530,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
 dependencies = [
  "glob",
  "libc",
- "libloading",
+ "libloading 0.8.6",
 ]
 
 [[package]]
@@ -1592,6 +1665,12 @@ dependencies = [
  "tiny-keccak",
 ]
 
+[[package]]
+name = "const_panic"
+version = "0.2.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2459fc9262a1aa204eb4b5764ad4f189caec88aea9634389c0a25f8be7f6265e"
+
 [[package]]
 name = "constant_time_eq"
 version = "0.3.1"
@@ -1633,6 +1712,21 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "core_extensions"
+version = "1.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92c71dc07c9721607e7a16108336048ee978c3a8b129294534272e8bac96c0ee"
+dependencies = [
+ "core_extensions_proc_macros",
+]
+
+[[package]]
+name = "core_extensions_proc_macros"
+version = "1.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69f3b219d28b6e3b4ac87bc1fc522e0803ab22e055da177bff0068c4150c61a6"
+
 [[package]]
 name = "cpufeatures"
 version = "0.2.17"
@@ -2187,6 +2281,26 @@ dependencies = [
  "paste",
 ]
 
+[[package]]
+name = "datafusion-ffi"
+version = "47.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5cf3fe9ab492c56daeb7beed526690d33622d388b8870472e0b7b7f55490338c"
+dependencies = [
+ "abi_stable",
+ "arrow",
+ "arrow-schema",
+ "async-ffi",
+ "async-trait",
+ "datafusion",
+ "datafusion-proto",
+ "futures",
+ "log",
+ "prost",
+ "semver",
+ "tokio",
+]
+
 [[package]]
 name = "datafusion-functions"
 version = "47.0.0"
@@ -2429,6 +2543,33 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "datafusion-proto"
+version = "47.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4a1afb2bdb05de7ff65be6883ebfd4ec027bd9f1f21c46aa3afd01927160a83"
+dependencies = [
+ "arrow",
+ "chrono",
+ "datafusion",
+ "datafusion-common",
+ "datafusion-expr",
+ "datafusion-proto-common",
+ "object_store",
+ "prost",
+]
+
+[[package]]
+name = "datafusion-proto-common"
+version = "47.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35b7a5876ebd6b564fb9a1fd2c3a2a9686b787071a256b47e4708f0916f9e46f"
+dependencies = [
+ "arrow",
+ "datafusion-common",
+ "prost",
+]
+
 [[package]]
 name = "datafusion-session"
 version = "47.0.0"
@@ -3021,6 +3162,15 @@ dependencies = [
  "slab",
 ]
 
+[[package]]
+name = "generational-arena"
+version = "0.2.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "generator"
 version = "0.8.4"
@@ -3936,6 +4086,12 @@ dependencies = [
  "web-time",
 ]
 
+[[package]]
+name = "indoc"
+version = "2.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd"
+
 [[package]]
 name = "inout"
 version = "0.1.3"
@@ -4183,6 +4339,16 @@ dependencies = [
  "rle-decode-fast",
 ]
 
+[[package]]
+name = "libloading"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f"
+dependencies = [
+ "cfg-if",
+ "winapi",
+]
+
 [[package]]
 name = "libloading"
 version = "0.8.6"
@@ -5278,6 +5444,29 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "prost"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5"
+dependencies = [
+ "bytes",
+ "prost-derive",
+]
+
+[[package]]
+name = "prost-derive"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d"
+dependencies = [
+ "anyhow",
+ "itertools 0.14.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.101",
+]
+
 [[package]]
 name = "psm"
 version = "0.1.25"
@@ -5327,6 +5516,81 @@ dependencies = [
  "syn 2.0.101",
 ]
 
+[[package]]
+name = "pyiceberg_core_rust"
+version = "0.5.1"
+dependencies = [
+ "arrow",
+ "datafusion-ffi",
+ "iceberg",
+ "iceberg-datafusion",
+ "pyo3",
+ "tokio",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5203598f366b11a02b13aa20cab591229ff0a89fd121a308a5df751d5fc9219"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "99636d423fa2ca130fa5acde3059308006d46f98caac629418e53f7ebb1e9999"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78f9cf92ba9c409279bc3305b5409d90db2d2c22392d443a87df3a1adad59e33"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b999cb1a6ce21f9a6b147dcf1be9ffedf02e0043aec74dc390f3007047cecd9"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn 2.0.101",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "822ece1c7e1012745607d5cf0bcb2874769f0f7cb34c4cde03b9358eb9ef911a"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn 2.0.101",
+]
+
 [[package]]
 name = "quad-rand"
 version = "0.2.3"
@@ -5637,6 +5901,15 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a35e8a6bf28cd121053a66aa2e6a2e3eaffad4a60012179f0e864aa5ffeff215"
 
+[[package]]
+name = "repr_offset"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb1070755bd29dffc19d0971cab794e607839ba2ef4b69a9e6fbc8733c1b72ea"
+dependencies = [
+ "tstr",
+]
+
 [[package]]
 name = "reqsign"
 version = "0.16.1"
@@ -6146,9 +6419,9 @@ dependencies = [
 
 [[package]]
 name = "semver"
-version = "1.0.25"
+version = "1.0.26"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f79dfe2d285b0488816f30e700a7438c5a73d816b5b7d3ac72fbc48b0d185e03"
+checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0"
 
 [[package]]
 name = "seq-macro"
@@ -6811,6 +7084,12 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
 
+[[package]]
+name = "target-lexicon"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a"
+
 [[package]]
 name = "tempfile"
 version = "3.20.0"
@@ -7191,6 +7470,21 @@ version = "0.2.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
 
+[[package]]
+name = "tstr"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f8e0294f14baae476d0dd0a2d780b2e24d66e349a9de876f5126777a37bdba7"
+dependencies = [
+ "tstr_proc_macros",
+]
+
+[[package]]
+name = "tstr_proc_macros"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e78122066b0cb818b8afd08f7ed22f7fdbc3e90815035726f0840d0d26c0747a"
+
 [[package]]
 name = "twox-hash"
 version = "1.6.3"
@@ -7207,6 +7501,12 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e7b17f197b3050ba473acf9181f7b1d3b66d1cf7356c6cc57886662276e65908"
 
+[[package]]
+name = "typed-arena"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a"
+
 [[package]]
 name = "typed-builder"
 version = "0.19.1"
@@ -7354,6 +7654,12 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"
 
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"
diff --git a/Cargo.toml b/Cargo.toml
index 3d12fc877..3f87a8571 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,7 +16,6 @@
 # under the License.
 
 [workspace]
-exclude = ["bindings/python"]
 members = [
   "crates/catalog/*",
   "crates/examples",
@@ -25,6 +24,7 @@ members = [
   "crates/integrations/*",
   "crates/sqllogictest",
   "crates/test_utils",
+  "bindings/python"
 ]
 resolver = "2"
 
diff --git a/bindings/python/src/runtime.rs b/bindings/python/src/runtime.rs
index 81d5b6dd6..35d757722 100644
--- a/bindings/python/src/runtime.rs
+++ b/bindings/python/src/runtime.rs
@@ -19,14 +19,25 @@ use std::sync::OnceLock;
 
 use tokio::runtime::{Handle, Runtime};
 
-static RUNTIME: OnceLock<Runtime> = OnceLock::new();
-
 pub fn runtime() -> Handle {
-    match Handle::try_current() {
-        Ok(h) => h.clone(),
-        _ => {
-            let rt = RUNTIME.get_or_init(|| Runtime::new().unwrap());
-            rt.handle().clone()
+    if let Ok(h) = Handle::try_current() {
+        h
+    } else {
+        static RUNTIME: OnceLock<Runtime> = OnceLock::new();
+        static PID: OnceLock<u32> = OnceLock::new();
+        let pid = std::process::id();
+        let runtime_pid = *PID.get_or_init(|| pid);
+        if pid != runtime_pid {
+            panic!(
+                "Forked process detected - current PID is {pid} but the tokio runtime was created by {runtime_pid}. 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.",
+            );
         }
+        RUNTIME
+            .get_or_init(|| Runtime::new().expect("Failed to create a tokio runtime."))
+            .handle()
+            .clone()
     }
 }