From 04a0a091786e287af774544e2f3173de9d7c750e Mon Sep 17 00:00:00 2001 From: Dino Viehland Date: Fri, 14 Feb 2025 12:00:46 -0800 Subject: [PATCH 1/3] SET_ADD should not lock --- Include/internal/pycore_setobject.h | 7 ++++ Objects/setobject.c | 59 +++++++++++++++++++++++++---- Python/bytecodes.c | 20 ++-------- Python/executor_cases.c.h | 51 +++++-------------------- Python/generated_cases.c.h | 51 ++++--------------------- 5 files changed, 80 insertions(+), 108 deletions(-) diff --git a/Include/internal/pycore_setobject.h b/Include/internal/pycore_setobject.h index 0494c07fe1869d..a65d17e6798a04 100644 --- a/Include/internal/pycore_setobject.h +++ b/Include/internal/pycore_setobject.h @@ -33,6 +33,13 @@ PyAPI_FUNC(int) _PySet_Contains(PySetObject *so, PyObject *key); // Clears the set without acquiring locks. Used by _PyCode_Fini. extern void _PySet_ClearInternal(PySetObject *so); +PyAPI_FUNC(int) _PySet_AddTakeRef(PySetObject *so, PyObject *key); + +PyAPI_FUNC(PyObject *) +_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n); + + + #ifdef __cplusplus } #endif diff --git a/Objects/setobject.c b/Objects/setobject.c index 9181bc27ac4e97..a89947c4281218 100644 --- a/Objects/setobject.c +++ b/Objects/setobject.c @@ -122,7 +122,7 @@ set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash) static int set_table_resize(PySetObject *, Py_ssize_t); static int -set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash) +set_add_entry_takeref(PySetObject *so, PyObject *key, Py_hash_t hash) { setentry *table; setentry *freeslot; @@ -133,12 +133,6 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash) int probes; int cmp; - _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so); - - /* Pre-increment is necessary to prevent arbitrary code in the rich - comparison from deallocating the key just before the insertion. */ - Py_INCREF(key); - restart: mask = so->mask; @@ -209,6 +203,57 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash) return -1; } +static int +set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash) +{ + _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so); + + /* Pre-increment is necessary to prevent arbitrary code in the rich + comparison from deallocating the key just before the insertion. */ + Py_INCREF(key); + + return set_add_entry_takeref(so, key, hash); +} + +int +_PySet_AddTakeRef(PySetObject *so, PyObject *key) +{ + Py_hash_t hash = _PyObject_HashFast(key); + if (hash == -1) { + return -1; + } + // We don't pre-increment here, the caller holds a strong + // reference to the object which we are stealing. + return set_add_entry_takeref(so, key, hash); +} + +PyObject * +_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n) +{ + PySetObject *set = (PySetObject *)PySet_New(NULL); + if (n == 0) { + return (PyObject *)set; + } + + if (set_table_resize(set, n*2) != 0) { + return NULL; + } + + int err = 0; + for (Py_ssize_t i = 0; i < n; i++) { + if (err == 0) { + err = _PySet_AddTakeRef(set, PyStackRef_AsPyObjectSteal(src[i])); + } else { + PyStackRef_CLOSE(src[i]); + } + } + if (err) { + Py_CLEAR(set); + } + + return (PyObject *)set; +} + /* Internal routine used by set_table_resize() to insert an item which is known to be absent from the set. Besides the performance benefit, diff --git a/Python/bytecodes.c b/Python/bytecodes.c index fe1465b36c7d04..8e592999cde11f 100644 --- a/Python/bytecodes.c +++ b/Python/bytecodes.c @@ -975,9 +975,8 @@ dummy_func( } inst(SET_ADD, (set, unused[oparg-1], v -- set, unused[oparg-1])) { - int err = PySet_Add(PyStackRef_AsPyObjectBorrow(set), - PyStackRef_AsPyObjectBorrow(v)); - PyStackRef_CLOSE(v); + int err = _PySet_AddTakeRef((PySetObject *)PyStackRef_AsPyObjectBorrow(set), + PyStackRef_AsPyObjectSteal(v)); ERROR_IF(err, error); } @@ -1911,22 +1910,11 @@ dummy_func( } inst(BUILD_SET, (values[oparg] -- set)) { - PyObject *set_o = PySet_New(NULL); + PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); if (set_o == NULL) { - DECREF_INPUTS(); - ERROR_IF(true, error); - } - int err = 0; - for (int i = 0; i < oparg; i++) { - if (err == 0) { - err = PySet_Add(set_o, PyStackRef_AsPyObjectBorrow(values[i])); - } - } - DECREF_INPUTS(); - if (err != 0) { - Py_DECREF(set_o); ERROR_IF(true, error); } + INPUTS_DEAD(); set = PyStackRef_FromPyObjectStealMortal(set_o); } diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h index 2bd009e37556fc..a31c46e02f70b9 100644 --- a/Python/executor_cases.c.h +++ b/Python/executor_cases.c.h @@ -1481,17 +1481,16 @@ v = stack_pointer[-1]; set = stack_pointer[-2 - (oparg-1)]; _PyFrame_SetStackPointer(frame, stack_pointer); - int err = PySet_Add(PyStackRef_AsPyObjectBorrow(set), - PyStackRef_AsPyObjectBorrow(v)); - stack_pointer = _PyFrame_GetStackPointer(frame); - stack_pointer += -1; - assert(WITHIN_STACK_BOUNDS()); - _PyFrame_SetStackPointer(frame, stack_pointer); - PyStackRef_CLOSE(v); + int err = _PySet_AddTakeRef((PySetObject *)PyStackRef_AsPyObjectBorrow(set), + PyStackRef_AsPyObjectSteal(v)); stack_pointer = _PyFrame_GetStackPointer(frame); if (err) { + stack_pointer += -1; + assert(WITHIN_STACK_BOUNDS()); JUMP_TO_ERROR(); } + stack_pointer += -1; + assert(WITHIN_STACK_BOUNDS()); break; } @@ -2671,48 +2670,16 @@ oparg = CURRENT_OPARG(); values = &stack_pointer[-oparg]; _PyFrame_SetStackPointer(frame, stack_pointer); - PyObject *set_o = PySet_New(NULL); + PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); stack_pointer = _PyFrame_GetStackPointer(frame); if (set_o == NULL) { - _PyFrame_SetStackPointer(frame, stack_pointer); - _PyStackRef tmp; - for (int _i = oparg; --_i >= 0;) { - tmp = values[_i]; - values[_i] = PyStackRef_NULL; - PyStackRef_CLOSE(tmp); - } - stack_pointer = _PyFrame_GetStackPointer(frame); stack_pointer += -oparg; assert(WITHIN_STACK_BOUNDS()); JUMP_TO_ERROR(); } - int err = 0; - for (int i = 0; i < oparg; i++) { - if (err == 0) { - _PyFrame_SetStackPointer(frame, stack_pointer); - err = PySet_Add(set_o, PyStackRef_AsPyObjectBorrow(values[i])); - stack_pointer = _PyFrame_GetStackPointer(frame); - } - } - _PyFrame_SetStackPointer(frame, stack_pointer); - _PyStackRef tmp; - for (int _i = oparg; --_i >= 0;) { - tmp = values[_i]; - values[_i] = PyStackRef_NULL; - PyStackRef_CLOSE(tmp); - } - stack_pointer = _PyFrame_GetStackPointer(frame); - stack_pointer += -oparg; - assert(WITHIN_STACK_BOUNDS()); - if (err != 0) { - _PyFrame_SetStackPointer(frame, stack_pointer); - Py_DECREF(set_o); - stack_pointer = _PyFrame_GetStackPointer(frame); - JUMP_TO_ERROR(); - } set = PyStackRef_FromPyObjectStealMortal(set_o); - stack_pointer[0] = set; - stack_pointer += 1; + stack_pointer[-oparg] = set; + stack_pointer += 1 - oparg; assert(WITHIN_STACK_BOUNDS()); break; } diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h index 9ce2b633d5e506..7be58f54597d28 100644 --- a/Python/generated_cases.c.h +++ b/Python/generated_cases.c.h @@ -1094,48 +1094,16 @@ _PyStackRef set; values = &stack_pointer[-oparg]; _PyFrame_SetStackPointer(frame, stack_pointer); - PyObject *set_o = PySet_New(NULL); + PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); stack_pointer = _PyFrame_GetStackPointer(frame); if (set_o == NULL) { - _PyFrame_SetStackPointer(frame, stack_pointer); - _PyStackRef tmp; - for (int _i = oparg; --_i >= 0;) { - tmp = values[_i]; - values[_i] = PyStackRef_NULL; - PyStackRef_CLOSE(tmp); - } - stack_pointer = _PyFrame_GetStackPointer(frame); stack_pointer += -oparg; assert(WITHIN_STACK_BOUNDS()); JUMP_TO_LABEL(error); } - int err = 0; - for (int i = 0; i < oparg; i++) { - if (err == 0) { - _PyFrame_SetStackPointer(frame, stack_pointer); - err = PySet_Add(set_o, PyStackRef_AsPyObjectBorrow(values[i])); - stack_pointer = _PyFrame_GetStackPointer(frame); - } - } - _PyFrame_SetStackPointer(frame, stack_pointer); - _PyStackRef tmp; - for (int _i = oparg; --_i >= 0;) { - tmp = values[_i]; - values[_i] = PyStackRef_NULL; - PyStackRef_CLOSE(tmp); - } - stack_pointer = _PyFrame_GetStackPointer(frame); - stack_pointer += -oparg; - assert(WITHIN_STACK_BOUNDS()); - if (err != 0) { - _PyFrame_SetStackPointer(frame, stack_pointer); - Py_DECREF(set_o); - stack_pointer = _PyFrame_GetStackPointer(frame); - JUMP_TO_LABEL(error); - } set = PyStackRef_FromPyObjectStealMortal(set_o); - stack_pointer[0] = set; - stack_pointer += 1; + stack_pointer[-oparg] = set; + stack_pointer += 1 - oparg; assert(WITHIN_STACK_BOUNDS()); DISPATCH(); } @@ -10641,17 +10609,14 @@ v = stack_pointer[-1]; set = stack_pointer[-2 - (oparg-1)]; _PyFrame_SetStackPointer(frame, stack_pointer); - int err = PySet_Add(PyStackRef_AsPyObjectBorrow(set), - PyStackRef_AsPyObjectBorrow(v)); - stack_pointer = _PyFrame_GetStackPointer(frame); - stack_pointer += -1; - assert(WITHIN_STACK_BOUNDS()); - _PyFrame_SetStackPointer(frame, stack_pointer); - PyStackRef_CLOSE(v); + int err = _PySet_AddTakeRef((PySetObject *)PyStackRef_AsPyObjectBorrow(set), + PyStackRef_AsPyObjectSteal(v)); stack_pointer = _PyFrame_GetStackPointer(frame); if (err) { - JUMP_TO_LABEL(error); + JUMP_TO_LABEL(pop_1_error); } + stack_pointer += -1; + assert(WITHIN_STACK_BOUNDS()); DISPATCH(); } From 4665a4c9e8dad33eb8a05771df872b078e4a3a78 Mon Sep 17 00:00:00 2001 From: Dino Viehland Date: Wed, 19 Feb 2025 14:11:02 -0800 Subject: [PATCH 2/3] Don't use helper that processes entire stack --- Include/internal/pycore_setobject.h | 3 --- Objects/setobject.c | 34 ++--------------------------- Python/bytecodes.c | 20 ++++++++++++++++- Python/executor_cases.c.h | 33 +++++++++++++++++++++++++++- Python/generated_cases.c.h | 33 +++++++++++++++++++++++++++- 5 files changed, 85 insertions(+), 38 deletions(-) diff --git a/Include/internal/pycore_setobject.h b/Include/internal/pycore_setobject.h index a65d17e6798a04..96f28dd9261c78 100644 --- a/Include/internal/pycore_setobject.h +++ b/Include/internal/pycore_setobject.h @@ -35,9 +35,6 @@ extern void _PySet_ClearInternal(PySetObject *so); PyAPI_FUNC(int) _PySet_AddTakeRef(PySetObject *so, PyObject *key); -PyAPI_FUNC(PyObject *) -_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n); - #ifdef __cplusplus diff --git a/Objects/setobject.c b/Objects/setobject.c index a89947c4281218..5ad83c3f1633b9 100644 --- a/Objects/setobject.c +++ b/Objects/setobject.c @@ -208,11 +208,7 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash) { _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so); - /* Pre-increment is necessary to prevent arbitrary code in the rich - comparison from deallocating the key just before the insertion. */ - Py_INCREF(key); - - return set_add_entry_takeref(so, key, hash); + return set_add_entry_takeref(so, Py_NewRef(key), hash); } int @@ -220,6 +216,7 @@ _PySet_AddTakeRef(PySetObject *so, PyObject *key) { Py_hash_t hash = _PyObject_HashFast(key); if (hash == -1) { + Py_DECREF(key); return -1; } // We don't pre-increment here, the caller holds a strong @@ -227,33 +224,6 @@ _PySet_AddTakeRef(PySetObject *so, PyObject *key) return set_add_entry_takeref(so, key, hash); } -PyObject * -_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n) -{ - PySetObject *set = (PySetObject *)PySet_New(NULL); - if (n == 0) { - return (PyObject *)set; - } - - if (set_table_resize(set, n*2) != 0) { - return NULL; - } - - int err = 0; - for (Py_ssize_t i = 0; i < n; i++) { - if (err == 0) { - err = _PySet_AddTakeRef(set, PyStackRef_AsPyObjectSteal(src[i])); - } else { - PyStackRef_CLOSE(src[i]); - } - } - if (err) { - Py_CLEAR(set); - } - - return (PyObject *)set; -} - /* Internal routine used by set_table_resize() to insert an item which is known to be absent from the set. Besides the performance benefit, diff --git a/Python/bytecodes.c b/Python/bytecodes.c index 8e592999cde11f..ebd133633e55ec 100644 --- a/Python/bytecodes.c +++ b/Python/bytecodes.c @@ -1910,10 +1910,28 @@ dummy_func( } inst(BUILD_SET, (values[oparg] -- set)) { - PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); + PyObject *set_o = PySet_New(NULL); if (set_o == NULL) { + DECREF_INPUTS(); ERROR_IF(true, error); } + + int err = 0; + for (Py_ssize_t i = 0; i < oparg; i++) { + _PyStackRef value = values[i]; + values[i] = PyStackRef_NULL; + if (err == 0) { + err = _PySet_AddTakeRef((PySetObject *)set_o, PyStackRef_AsPyObjectSteal(value)); + } + else { + PyStackRef_CLOSE(value); + } + } + if (err) { + Py_DECREF(set_o); + ERROR_IF(true, error); + } + INPUTS_DEAD(); set = PyStackRef_FromPyObjectStealMortal(set_o); } diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h index a31c46e02f70b9..2942680a222fb2 100644 --- a/Python/executor_cases.c.h +++ b/Python/executor_cases.c.h @@ -2670,9 +2670,40 @@ oparg = CURRENT_OPARG(); values = &stack_pointer[-oparg]; _PyFrame_SetStackPointer(frame, stack_pointer); - PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); + PyObject *set_o = PySet_New(NULL); stack_pointer = _PyFrame_GetStackPointer(frame); if (set_o == NULL) { + _PyFrame_SetStackPointer(frame, stack_pointer); + _PyStackRef tmp; + for (int _i = oparg; --_i >= 0;) { + tmp = values[_i]; + values[_i] = PyStackRef_NULL; + PyStackRef_CLOSE(tmp); + } + stack_pointer = _PyFrame_GetStackPointer(frame); + stack_pointer += -oparg; + assert(WITHIN_STACK_BOUNDS()); + JUMP_TO_ERROR(); + } + int err = 0; + for (Py_ssize_t i = 0; i < oparg; i++) { + _PyStackRef value = values[i]; + values[i] = PyStackRef_NULL; + if (err == 0) { + _PyFrame_SetStackPointer(frame, stack_pointer); + err = _PySet_AddTakeRef((PySetObject *)set_o, PyStackRef_AsPyObjectSteal(value)); + stack_pointer = _PyFrame_GetStackPointer(frame); + } + else { + _PyFrame_SetStackPointer(frame, stack_pointer); + PyStackRef_CLOSE(value); + stack_pointer = _PyFrame_GetStackPointer(frame); + } + } + if (err) { + _PyFrame_SetStackPointer(frame, stack_pointer); + Py_DECREF(set_o); + stack_pointer = _PyFrame_GetStackPointer(frame); stack_pointer += -oparg; assert(WITHIN_STACK_BOUNDS()); JUMP_TO_ERROR(); diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h index 7be58f54597d28..a2e6c474d29204 100644 --- a/Python/generated_cases.c.h +++ b/Python/generated_cases.c.h @@ -1094,9 +1094,40 @@ _PyStackRef set; values = &stack_pointer[-oparg]; _PyFrame_SetStackPointer(frame, stack_pointer); - PyObject *set_o = _PySet_FromStackRefSteal(values, oparg); + PyObject *set_o = PySet_New(NULL); stack_pointer = _PyFrame_GetStackPointer(frame); if (set_o == NULL) { + _PyFrame_SetStackPointer(frame, stack_pointer); + _PyStackRef tmp; + for (int _i = oparg; --_i >= 0;) { + tmp = values[_i]; + values[_i] = PyStackRef_NULL; + PyStackRef_CLOSE(tmp); + } + stack_pointer = _PyFrame_GetStackPointer(frame); + stack_pointer += -oparg; + assert(WITHIN_STACK_BOUNDS()); + JUMP_TO_LABEL(error); + } + int err = 0; + for (Py_ssize_t i = 0; i < oparg; i++) { + _PyStackRef value = values[i]; + values[i] = PyStackRef_NULL; + if (err == 0) { + _PyFrame_SetStackPointer(frame, stack_pointer); + err = _PySet_AddTakeRef((PySetObject *)set_o, PyStackRef_AsPyObjectSteal(value)); + stack_pointer = _PyFrame_GetStackPointer(frame); + } + else { + _PyFrame_SetStackPointer(frame, stack_pointer); + PyStackRef_CLOSE(value); + stack_pointer = _PyFrame_GetStackPointer(frame); + } + } + if (err) { + _PyFrame_SetStackPointer(frame, stack_pointer); + Py_DECREF(set_o); + stack_pointer = _PyFrame_GetStackPointer(frame); stack_pointer += -oparg; assert(WITHIN_STACK_BOUNDS()); JUMP_TO_LABEL(error); From 1ff9eb34bc85451b5d212051efc4d61ef9355603 Mon Sep 17 00:00:00 2001 From: Dino Viehland Date: Wed, 19 Mar 2025 11:17:57 -0700 Subject: [PATCH 3/3] Fix formatting --- Include/internal/pycore_setobject.h | 2 -- Python/bytecodes.c | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/Include/internal/pycore_setobject.h b/Include/internal/pycore_setobject.h index 96f28dd9261c78..24d0135ed1aeca 100644 --- a/Include/internal/pycore_setobject.h +++ b/Include/internal/pycore_setobject.h @@ -35,8 +35,6 @@ extern void _PySet_ClearInternal(PySetObject *so); PyAPI_FUNC(int) _PySet_AddTakeRef(PySetObject *so, PyObject *key); - - #ifdef __cplusplus } #endif diff --git a/Python/bytecodes.c b/Python/bytecodes.c index ebd133633e55ec..0db4da013ff40a 100644 --- a/Python/bytecodes.c +++ b/Python/bytecodes.c @@ -976,7 +976,7 @@ dummy_func( inst(SET_ADD, (set, unused[oparg-1], v -- set, unused[oparg-1])) { int err = _PySet_AddTakeRef((PySetObject *)PyStackRef_AsPyObjectBorrow(set), - PyStackRef_AsPyObjectSteal(v)); + PyStackRef_AsPyObjectSteal(v)); ERROR_IF(err, error); }