more threading fix

youknowone · youknowone · commit 800dcd476e3c · 2026-01-10T20:50:52.000+09:00
diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
@@ -343,6 +343,8 @@ def run(self):
             t.join()
         # else the thread is still running, and we have no way to kill it
 
+    # TODO: RUSTPYTHON: threading._start_new_thread is not exposed
+    @unittest.skip("TODO: RUSTPYTHON; threading._start_new_thread not exposed")
     def test_limbo_cleanup(self):
         # Issue 7481: Failure to start thread should cleanup the limbo map.
         def fail_new_thread(*args):
@@ -742,6 +744,8 @@ def f():
         rc, out, err = assert_python_ok("-c", code)
         self.assertEqual(err, b"")
 
+    # TODO: RUSTPYTHON: Thread._tstate_lock is a CPython implementation detail
+    @unittest.skip("TODO: RUSTPYTHON; Thread._tstate_lock not implemented")
     def test_tstate_lock(self):
         # Test an implementation detail of Thread objects.
         started = _thread.allocate_lock()
diff --git a/crates/vm/src/stdlib/thread.rs b/crates/vm/src/stdlib/thread.rs
@@ -19,7 +19,6 @@ pub(crate) mod _thread {
         lock_api::{RawMutex as RawMutexT, RawMutexTimed, RawReentrantMutex},
     };
     use std::thread;
-    use thread_local::ThreadLocal;
 
     // PYTHREAD_NAME: show current thread name
     pub const PYTHREAD_NAME: Option<&str> = {
@@ -575,23 +574,77 @@ pub(crate) mod _thread {
         Ok(())
     }
 
+    // Thread-local storage for cleanup guards
+    // When a thread terminates, the guard is dropped, which triggers cleanup
+    thread_local! {
+        static LOCAL_GUARDS: std::cell::RefCell<Vec<LocalGuard>> = const { std::cell::RefCell::new(Vec::new()) };
+    }
+
+    // Guard that removes thread-local data when dropped
+    struct LocalGuard {
+        local: std::sync::Weak<LocalData>,
+        thread_id: std::thread::ThreadId,
+    }
+
+    impl Drop for LocalGuard {
+        fn drop(&mut self) {
+            // eprintln!("[DEBUG] LocalGuard::drop called for thread {:?}", self.thread_id);
+            if let Some(local_data) = self.local.upgrade() {
+                local_data.data.lock().remove(&self.thread_id);
+            }
+        }
+    }
+
+    // Shared data structure for Local
+    struct LocalData {
+        data: parking_lot::Mutex<std::collections::HashMap<std::thread::ThreadId, PyDictRef>>,
+    }
+
+    impl std::fmt::Debug for LocalData {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.debug_struct("LocalData").finish_non_exhaustive()
+        }
+    }
+
     #[pyattr]
     #[pyclass(module = "thread", name = "_local")]
     #[derive(Debug, PyPayload)]
     struct Local {
-        data: ThreadLocal<PyDictRef>,
+        inner: std::sync::Arc<LocalData>,
     }
 
     #[pyclass(with(GetAttr, SetAttr), flags(BASETYPE))]
     impl Local {
         fn l_dict(&self, vm: &VirtualMachine) -> PyDictRef {
-            self.data.get_or(|| vm.ctx.new_dict()).clone()
+            let thread_id = std::thread::current().id();
+            let mut data = self.inner.data.lock();
+
+            if let Some(dict) = data.get(&thread_id) {
+                return dict.clone();
+            }
+
+            // Create new dict for this thread
+            let dict = vm.ctx.new_dict();
+            data.insert(thread_id, dict.clone());
+
+            // Register cleanup guard for this thread
+            let guard = LocalGuard {
+                local: std::sync::Arc::downgrade(&self.inner),
+                thread_id,
+            };
+            LOCAL_GUARDS.with(|guards| {
+                guards.borrow_mut().push(guard);
+            });
+
+            dict
         }
 
         #[pyslot]
         fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
             Self {
-                data: ThreadLocal::new(),
+                inner: std::sync::Arc::new(LocalData {
+                    data: parking_lot::Mutex::new(std::collections::HashMap::new()),
+                }),
             }
             .into_ref_with_type(vm, cls)
             .map(Into::into)
@@ -706,8 +759,8 @@ pub(crate) mod _thread {
         }
 
         #[pymethod]
-        fn join(&self, timeout: OptionalArg<f64>, vm: &VirtualMachine) -> PyResult<()> {
-            let _timeout_val = timeout.unwrap_or(-1.0);
+        fn join(&self, timeout: OptionalArg<Option<f64>>, vm: &VirtualMachine) -> PyResult<()> {
+            let _timeout_val = timeout.flatten().unwrap_or(-1.0);
 
             // Check for self-join and if already joined
             let join_handle = {
