From b6524b9295a1989b025f5a9cda7e76d53df7a35a Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Wed, 1 May 2024 17:53:49 +0900
Subject: [PATCH] sys.get_coroutine_origin_tracking_depth

---
 vm/src/stdlib/sys.rs | 14 ++++++++++++++
 vm/src/vm/thread.rs  |  9 +++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/vm/src/stdlib/sys.rs b/vm/src/stdlib/sys.rs
index 55e9c12b7e..d6b10f4d2d 100644
--- a/vm/src/stdlib/sys.rs
+++ b/vm/src/stdlib/sys.rs
@@ -692,6 +692,20 @@ mod sys {
         vm.use_tracing.set(tracing);
     }
 
+    #[pyfunction]
+    fn set_coroutine_origin_tracking_depth(depth: i32, vm: &VirtualMachine) -> PyResult<()> {
+        if depth < 0 {
+            return Err(vm.new_value_error("depth must be >= 0".to_owned()));
+        }
+        crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.set(depth as _));
+        Ok(())
+    }
+
+    #[pyfunction]
+    fn get_coroutine_origin_tracking_depth() -> i32 {
+        crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.get()) as _
+    }
+
     /// sys.flags
     ///
     /// Flags provided through command line arguments or environment vars.
diff --git a/vm/src/vm/thread.rs b/vm/src/vm/thread.rs
index 4460c39cd7..82110f73da 100644
--- a/vm/src/vm/thread.rs
+++ b/vm/src/vm/thread.rs
@@ -1,14 +1,16 @@
 use crate::{AsObject, PyObject, VirtualMachine};
 use itertools::Itertools;
 use std::{
-    cell::RefCell,
-    ptr::{null, NonNull},
+    cell::{Cell, RefCell},
+    ptr::NonNull,
     thread_local,
 };
 
 thread_local! {
     pub(super) static VM_STACK: RefCell<Vec<NonNull<VirtualMachine>>> = Vec::with_capacity(1).into();
-    static VM_CURRENT: RefCell<*const VirtualMachine> = null::<VirtualMachine>().into();
+    static VM_CURRENT: RefCell<*const VirtualMachine> = std::ptr::null::<VirtualMachine>().into();
+
+    pub(crate) static COROUTINE_ORIGIN_TRACKING_DEPTH: Cell<u32> = const { Cell::new(0) };
 }
 
 pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
@@ -142,7 +144,6 @@ impl VirtualMachine {
     /// specific guaranteed behavior.
     #[cfg(feature = "threading")]
     pub fn new_thread(&self) -> ThreadedVirtualMachine {
-        use std::cell::Cell;
         let vm = VirtualMachine {
             builtins: self.builtins.clone(),
             sys_module: self.sys_module.clone(),