gh-103845: Remove line & instruction instrumentations before adding them back (GH-103851)

3 files changed, 70 insertions, 8 deletions

diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py
index 738ace9..a493bb5 100644
--- a/Lib/test/test_monitoring.py
+++ b/Lib/test/test_monitoring.py
@@ -876,6 +876,42 @@ class TestLineAndInstructionEvents(CheckEvents):
             ('instruction', 'func3', 34),
             ('line', 'check_events', 11)])
 
+    def test_with_restart(self):
+        def func1():
+            line1 = 1
+            line2 = 2
+            line3 = 3
+
+        self.check_events(func1, recorders = LINE_AND_INSTRUCTION_RECORDERS, expected = [
+            ('line', 'check_events', 10),
+            ('line', 'func1', 1),
+            ('instruction', 'func1', 2),
+            ('instruction', 'func1', 4),
+            ('line', 'func1', 2),
+            ('instruction', 'func1', 6),
+            ('instruction', 'func1', 8),
+            ('line', 'func1', 3),
+            ('instruction', 'func1', 10),
+            ('instruction', 'func1', 12),
+            ('instruction', 'func1', 14),
+            ('line', 'check_events', 11)])
+
+        sys.monitoring.restart_events()
+
+        self.check_events(func1, recorders = LINE_AND_INSTRUCTION_RECORDERS, expected = [
+            ('line', 'check_events', 10),
+            ('line', 'func1', 1),
+            ('instruction', 'func1', 2),
+            ('instruction', 'func1', 4),
+            ('line', 'func1', 2),
+            ('instruction', 'func1', 6),
+            ('instruction', 'func1', 8),
+            ('line', 'func1', 3),
+            ('instruction', 'func1', 10),
+            ('instruction', 'func1', 12),
+            ('instruction', 'func1', 14),
+            ('line', 'check_events', 11)])
+
 class TestInstallIncrementallly(MonitoringTestBase, unittest.TestCase):
 
     def check_events(self, func, must_include, tool=TEST_TOOL, recorders=(ExceptionRecorder,)):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-04-25-20-56-01.gh-issue-103845.V7NYFn.rst b/Misc/NEWS.d/next/Core and Builtins/2023-04-25-20-56-01.gh-issue-103845.V7NYFn.rst
new file mode 100644
index 0000000..e843485
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-04-25-20-56-01.gh-issue-103845.V7NYFn.rst
@@ -0,0 +1 @@
+Remove both line and instruction instrumentation before adding new ones for monitoring, to avoid newly added instrumentation being removed immediately.
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index c5bbbda..a142324 100644
--- a/Python/instrumentation.c
+++ b/Python/instrumentation.c
@@ -1477,25 +1477,25 @@ _Py_Instrument(PyCodeObject *code, PyInterpreterState *interp)
             }
         }
     }
-    uint8_t new_line_tools = new_events.tools[PY_MONITORING_EVENT_LINE];
+
+    // GH-103845: We need to remove both the line and instruction instrumentation before
+    // adding new ones, otherwise we may remove the newly added instrumentation.
+
     uint8_t removed_line_tools = removed_events.tools[PY_MONITORING_EVENT_LINE];
-    if (new_line_tools | removed_line_tools) {
+    uint8_t removed_per_instruction_tools = removed_events.tools[PY_MONITORING_EVENT_INSTRUCTION];
+
+    if (removed_line_tools) {
         _PyCoLineInstrumentationData *line_data = code->_co_monitoring->lines;
         for (int i = code->_co_firsttraceable; i < code_len;) {
             if (line_data[i].original_opcode) {
                 if (removed_line_tools) {
                     remove_line_tools(code, i, removed_line_tools);
                 }
-                if (new_line_tools) {
-                    add_line_tools(code, i, new_line_tools);
-                }
             }
             i += instruction_length(code, i);
         }
     }
-    uint8_t new_per_instruction_tools = new_events.tools[PY_MONITORING_EVENT_INSTRUCTION];
-    uint8_t removed_per_instruction_tools = removed_events.tools[PY_MONITORING_EVENT_INSTRUCTION];
-    if (new_per_instruction_tools | removed_per_instruction_tools) {
+    if (removed_per_instruction_tools) {
         for (int i = code->_co_firsttraceable; i < code_len;) {
             int opcode = _Py_GetBaseOpcode(code, i);
             if (opcode == RESUME || opcode == END_FOR) {
@@ -1505,6 +1505,31 @@ _Py_Instrument(PyCodeObject *code, PyInterpreterState *interp)
             if (removed_per_instruction_tools) {
                 remove_per_instruction_tools(code, i, removed_per_instruction_tools);
             }
+            i += instruction_length(code, i);
+        }
+    }
+
+    uint8_t new_line_tools = new_events.tools[PY_MONITORING_EVENT_LINE];
+    uint8_t new_per_instruction_tools = new_events.tools[PY_MONITORING_EVENT_INSTRUCTION];
+
+    if (new_line_tools) {
+        _PyCoLineInstrumentationData *line_data = code->_co_monitoring->lines;
+        for (int i = code->_co_firsttraceable; i < code_len;) {
+            if (line_data[i].original_opcode) {
+                if (new_line_tools) {
+                    add_line_tools(code, i, new_line_tools);
+                }
+            }
+            i += instruction_length(code, i);
+        }
+    }
+    if (new_per_instruction_tools) {
+        for (int i = code->_co_firsttraceable; i < code_len;) {
+            int opcode = _Py_GetBaseOpcode(code, i);
+            if (opcode == RESUME || opcode == END_FOR) {
+                i += instruction_length(code, i);
+                continue;
+            }
             if (new_per_instruction_tools) {
                 add_per_instruction_tools(code, i, new_per_instruction_tools);
             }