more fixes

Author: kumaraditya303

Python/instrumentation.c

@@ -40,14 +40,9 @@
 
 #ifdef Py_GIL_DISABLED
 
-#define LOCK_CODE(code)                                             \
-    assert(!_PyInterpreterState_GET()->stoptheworld.world_stopped); \
-    Py_BEGIN_CRITICAL_SECTION(code)
-
-#define UNLOCK_CODE()   Py_END_CRITICAL_SECTION()
-
 #define MODIFY_BYTECODE(code, func, ...)                       \
     do {                                                       \
+        ASSERT_WORLD_STOPPED();                                \
         PyCodeObject *co = (code);                             \
         for (Py_ssize_t i = 0; i < code->co_tlbc->size; i++) { \
             char *bc = co->co_tlbc->entries[i];                \
@@ -60,8 +55,6 @@
 
 #else
 
-#define LOCK_CODE(code)
-#define UNLOCK_CODE()
 #define MODIFY_BYTECODE(code, func, ...) \
     (func)(code, _PyCode_CODE(code), __VA_ARGS__)
 
@@ -535,7 +528,7 @@ valid_opcode(int opcode)
 static void
 sanity_check_instrumentation(PyCodeObject *code)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
 
     _PyCoMonitoringData *data = code->_co_monitoring;
     if (data == NULL) {
@@ -945,7 +938,7 @@ add_per_instruction_tools(PyCodeObject * code, int offset, int tools)
 static void
 remove_per_instruction_tools(PyCodeObject * code, int offset, int tools)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
 
     _PyCoMonitoringData *monitoring = code->_co_monitoring;
     assert(code->_co_monitoring);
@@ -1060,15 +1053,15 @@ set_global_version(PyThreadState *tstate, uint32_t version)
 static bool
 is_version_up_to_date(PyCodeObject *code, PyInterpreterState *interp)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
     return global_version(interp) == code->_co_instrumentation_version;
 }
 
 #ifndef NDEBUG
 static bool
 instrumentation_cross_checks(PyInterpreterState *interp, PyCodeObject *code)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
     _Py_LocalMonitors expected = local_union(
         interp->monitors,
         code->_co_monitoring->local_monitors);
@@ -1079,10 +1072,10 @@ static int
 debug_check_sanity(PyInterpreterState *interp, PyCodeObject *code)
 {
     int res;
-    LOCK_CODE(code);
+    _PyEval_StopTheWorld(interp);
     res = is_version_up_to_date(code, interp) &&
           instrumentation_cross_checks(interp, code);
-    UNLOCK_CODE();
+    _PyEval_StartTheWorld(interp);
     return res;
 }
 
@@ -1381,9 +1374,9 @@ _Py_call_instrumentation_line(PyThreadState *tstate, _PyInterpreterFrame* frame,
         }
         else {
             /* DISABLE  */
-            LOCK_CODE(code);
+            _PyEval_StopTheWorld(tstate->interp);
             remove_line_tools(code, i, 1 << tool);
-            UNLOCK_CODE();
+            _PyEval_StartTheWorld(tstate->interp);
         }
     } while (tools);
     Py_DECREF(line_obj);
@@ -1438,9 +1431,9 @@ _Py_call_instrumentation_instruction(PyThreadState *tstate, _PyInterpreterFrame*
         }
         else {
             /* DISABLE  */
-            LOCK_CODE(code);
+            _PyEval_StopTheWorld(tstate->interp);
             remove_per_instruction_tools(code, offset, 1 << tool);
-            UNLOCK_CODE();
+            _PyEval_StartTheWorld(tstate->interp);
         }
     }
     Py_DECREF(offset_obj);
@@ -1451,7 +1444,7 @@ _Py_call_instrumentation_instruction(PyThreadState *tstate, _PyInterpreterFrame*
 static void
 initialize_tools(PyCodeObject *code)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
     uint8_t* tools = code->_co_monitoring->tools;
 
     assert(tools != NULL);
@@ -1507,7 +1500,7 @@ initialize_tools(PyCodeObject *code)
 static void
 initialize_lines(PyCodeObject *code, int bytes_per_entry)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
     _PyCoLineInstrumentationData *line_data = code->_co_monitoring->lines;
 
     assert(line_data != NULL);
@@ -1634,7 +1627,7 @@ initialize_lines(PyCodeObject *code, int bytes_per_entry)
 static void
 initialize_line_tools(PyCodeObject *code, _Py_LocalMonitors *all_events)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
     uint8_t *line_tools = code->_co_monitoring->line_tools;
 
     assert(line_tools != NULL);
@@ -1647,7 +1640,7 @@ initialize_line_tools(PyCodeObject *code, _Py_LocalMonitors *all_events)
 static int
 allocate_instrumentation_data(PyCodeObject *code)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
 
     if (code->_co_monitoring == NULL) {
         code->_co_monitoring = PyMem_Malloc(sizeof(_PyCoMonitoringData));
@@ -1669,7 +1662,7 @@ allocate_instrumentation_data(PyCodeObject *code)
 static int
 update_instrumentation_data(PyCodeObject *code, PyInterpreterState *interp)
 {
-    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    ASSERT_WORLD_STOPPED();
 
     int code_len = (int)Py_SIZE(code);
     if (allocate_instrumentation_data(code)) {
@@ -2995,9 +2988,10 @@ branch_handler_vectorcall(
             // Orphaned NOT_TAKEN -- Jump removed by the compiler
             return res;
         }
-        LOCK_CODE(code);
+        PyInterpreterState *interp = _PyInterpreterState_GET();
+        _PyEval_StopTheWorld(interp);
         remove_tools(code, offset, other_event, 1 << self->tool_id);
-        UNLOCK_CODE();
+        _PyEval_StartTheWorld(interp);
     }
     return res;
 }