gh-117657: Fix small issues with instrumentation and TSAN (#118064)

Small TSAN fixups for instrumentation

7 files changed, 25 insertions, 13 deletions

diff --git a/Include/internal/pycore_pyatomic_ft_wrappers.h b/Include/internal/pycore_pyatomic_ft_wrappers.h
index bbfc462..83f02c9 100644
--- a/Include/internal/pycore_pyatomic_ft_wrappers.h
+++ b/Include/internal/pycore_pyatomic_ft_wrappers.h
@@ -39,6 +39,10 @@ extern "C" {
     _Py_atomic_load_uint8(&value)
 #define FT_ATOMIC_STORE_UINT8(value, new_value) \
     _Py_atomic_store_uint8(&value, new_value)
+#define FT_ATOMIC_LOAD_UINT8_RELAXED(value) \
+    _Py_atomic_load_uint8_relaxed(&value)
+#define FT_ATOMIC_LOAD_UINT16_RELAXED(value) \
+    _Py_atomic_load_uint16_relaxed(&value)
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) \
     _Py_atomic_store_ptr_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) \
@@ -49,7 +53,8 @@ extern "C" {
     _Py_atomic_store_ssize_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) \
     _Py_atomic_store_uint8_relaxed(&value, new_value)
-
+#define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) \
+    _Py_atomic_store_uint16_relaxed(&value, new_value)
 #else
 #define FT_ATOMIC_LOAD_PTR(value) value
 #define FT_ATOMIC_STORE_PTR(value, new_value) value = new_value
@@ -62,12 +67,14 @@ extern "C" {
 #define FT_ATOMIC_LOAD_PTR_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT8(value) value
 #define FT_ATOMIC_STORE_UINT8(value, new_value) value = new_value
+#define FT_ATOMIC_LOAD_UINT8_RELAXED(value) value
+#define FT_ATOMIC_LOAD_UINT16_RELAXED(value) value
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINTPTR_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) value = new_value
-
+#define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) value = new_value
 #endif
 
 #ifdef __cplusplus
diff --git a/Objects/genobject.c b/Objects/genobject.c
index 04736a0..a1ed1cb 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -11,6 +11,7 @@
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_opcode_utils.h"  // RESUME_AFTER_YIELD_FROM
+#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_*
 #include "pycore_pyerrors.h"      // _PyErr_ClearExcState()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 
@@ -329,10 +330,11 @@ gen_close_iter(PyObject *yf)
 static inline bool
 is_resume(_Py_CODEUNIT *instr)
 {
+    uint8_t code = FT_ATOMIC_LOAD_UINT8_RELAXED(instr->op.code);
     return (
-        instr->op.code == RESUME ||
-        instr->op.code == RESUME_CHECK ||
-        instr->op.code == INSTRUMENTED_RESUME
+        code == RESUME ||
+        code == RESUME_CHECK ||
+        code == INSTRUMENTED_RESUME
     );
 }
 
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index eee8b32..5bb7e12 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -20,7 +20,7 @@
 #include "pycore_object.h"        // _PyObject_GC_TRACK()
 #include "pycore_opcode_metadata.h"  // uop names
 #include "pycore_opcode_utils.h"  // MAKE_FUNCTION_*
-#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_PTR_ACQUIRE
+#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_*
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException()
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
 #include "pycore_range.h"         // _PyRangeIterObject
@@ -163,7 +163,7 @@ dummy_func(
                 if ((oparg & RESUME_OPARG_LOCATION_MASK) < RESUME_AFTER_YIELD_FROM) {
                     CHECK_EVAL_BREAKER();
                 }
-                this_instr->op.code = RESUME_CHECK;
+                FT_ATOMIC_STORE_UINT8_RELAXED(this_instr->op.code, RESUME_CHECK);
             }
         }
 
diff --git a/Python/ceval.c b/Python/ceval.c
index d130c73..8b23bc6 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -20,7 +20,7 @@
 #include "pycore_opcode_metadata.h" // EXTRA_CASES
 #include "pycore_optimizer.h"     // _PyUOpExecutor_Type
 #include "pycore_opcode_utils.h"  // MAKE_FUNCTION_*
-#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_PTR_ACQUIRE
+#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_*
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException()
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
 #include "pycore_range.h"         // _PyRangeIterObject
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
index 1a8554a..c88a07c 100644
--- a/Python/ceval_macros.h
+++ b/Python/ceval_macros.h
@@ -162,7 +162,7 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 /* The integer overflow is checked by an assertion below. */
 #define INSTR_OFFSET() ((int)(next_instr - _PyCode_CODE(_PyFrame_GetCode(frame))))
 #define NEXTOPARG()  do { \
-        _Py_CODEUNIT word = *next_instr; \
+        _Py_CODEUNIT word  = {.cache = FT_ATOMIC_LOAD_UINT16_RELAXED(*(uint16_t*)next_instr)}; \
         opcode = word.op.code; \
         oparg = word.op.arg; \
     } while (0)
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 7c1cc14..a5bb293 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -4955,7 +4955,7 @@
                 if ((oparg & RESUME_OPARG_LOCATION_MASK) < RESUME_AFTER_YIELD_FROM) {
                     CHECK_EVAL_BREAKER();
                 }
-                this_instr->op.code = RESUME_CHECK;
+                FT_ATOMIC_STORE_UINT8_RELAXED(this_instr->op.code, RESUME_CHECK);
             }
             DISPATCH();
         }
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index 328a3b1..ce97c3a 100644
--- a/Python/instrumentation.c
+++ b/Python/instrumentation.c
@@ -626,9 +626,10 @@ de_instrument(PyCodeObject *code, int i, int event)
         return;
     }
     CHECK(_PyOpcode_Deopt[deinstrumented] == deinstrumented);
-    *opcode_ptr = deinstrumented;
+    FT_ATOMIC_STORE_UINT8_RELAXED(*opcode_ptr, deinstrumented);
     if (_PyOpcode_Caches[deinstrumented]) {
-        instr[1].counter = adaptive_counter_warmup();
+        FT_ATOMIC_STORE_UINT16_RELAXED(instr[1].counter.as_counter,
+                                       adaptive_counter_warmup().as_counter);
     }
 }
 
@@ -703,8 +704,10 @@ instrument(PyCodeObject *code, int i)
         int deopt = _PyOpcode_Deopt[opcode];
         int instrumented = INSTRUMENTED_OPCODES[deopt];
         assert(instrumented);
-        *opcode_ptr = instrumented;
+        FT_ATOMIC_STORE_UINT8_RELAXED(*opcode_ptr, instrumented);
         if (_PyOpcode_Caches[deopt]) {
+          FT_ATOMIC_STORE_UINT16_RELAXED(instr[1].counter.as_counter,
+                                         adaptive_counter_warmup().as_counter);
             instr[1].counter = adaptive_counter_warmup();
         }
     }