Skip to content

Commit

Permalink
pythongh-125245: Fix ZoneInfo data race in free threading build
Browse files Browse the repository at this point in the history
Lock `ZoneInfoType` to protect accesses to `ZONEINFO_STRONG_CACHE`.
Refactor the `tp_new` handler to use Argument Clinic so that we can just
use `@critical_section` annotations on the relevant functions.

Also use `PyDict_SetDefaultRef` instead of `PyDict_SetDefault` when
inserting into the `TIMEDELTA_CACHE`.
  • Loading branch information
colesbury committed Oct 10, 2024
1 parent dd0ee20 commit bcaa724
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 13 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix data race when creating :class:`ZoneInfo` objects in the free threading
build.
34 changes: 22 additions & 12 deletions Modules/_zoneinfo.c
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#endif

#include "Python.h"
#include "pycore_critical_section.h" // _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED()
#include "pycore_long.h" // _PyLong_GetOne()
#include "pycore_pyerrors.h" // _PyErr_ChainExceptions1()

Expand Down Expand Up @@ -298,15 +299,20 @@ get_weak_cache(zoneinfo_state *state, PyTypeObject *type)
}
}

/*[clinic input]
@critical_section
@classmethod
zoneinfo.ZoneInfo.__new__
key: object
Create a new ZoneInfo instance.
[clinic start generated code]*/

static PyObject *
zoneinfo_new(PyTypeObject *type, PyObject *args, PyObject *kw)
zoneinfo_ZoneInfo_impl(PyTypeObject *type, PyObject *key)
/*[clinic end generated code: output=95e61dab86bb95c3 input=4619eb0305327e83]*/
{
PyObject *key = NULL;
static char *kwlist[] = {"key", NULL};
if (PyArg_ParseTupleAndKeywords(args, kw, "O", kwlist, &key) == 0) {
return NULL;
}

zoneinfo_state *state = zoneinfo_get_state_by_self(type);
PyObject *instance = zone_from_strong_cache(state, type, key);
if (instance != NULL || PyErr_Occurred()) {
Expand Down Expand Up @@ -467,6 +473,7 @@ zoneinfo_ZoneInfo_no_cache_impl(PyTypeObject *type, PyTypeObject *cls,
}

/*[clinic input]
@critical_section
@classmethod
zoneinfo.ZoneInfo.clear_cache
Expand All @@ -481,7 +488,7 @@ Clear the ZoneInfo cache.
static PyObject *
zoneinfo_ZoneInfo_clear_cache_impl(PyTypeObject *type, PyTypeObject *cls,
PyObject *only_keys)
/*[clinic end generated code: output=114d9b7c8a22e660 input=e32ca3bb396788ba]*/
/*[clinic end generated code: output=114d9b7c8a22e660 input=35944715df26d24e]*/
{
zoneinfo_state *state = zoneinfo_get_state_by_cls(cls);
PyObject *weak_cache = get_weak_cache(state, type);
Expand Down Expand Up @@ -821,7 +828,7 @@ zoneinfo_ZoneInfo__unpickle_impl(PyTypeObject *type, PyTypeObject *cls,
return NULL;
}

PyObject *rv = zoneinfo_new(type, val_args, NULL);
PyObject *rv = zoneinfo_ZoneInfo(type, val_args, NULL);

Py_DECREF(val_args);
return rv;
Expand Down Expand Up @@ -858,8 +865,7 @@ load_timedelta(zoneinfo_state *state, long seconds)
0, seconds, 0, 1, PyDateTimeAPI->DeltaType);

if (tmp != NULL) {
rv = PyDict_SetDefault(state->TIMEDELTA_CACHE, pyoffset, tmp);
Py_XINCREF(rv);
PyDict_SetDefaultRef(state->TIMEDELTA_CACHE, pyoffset, tmp, &rv);
Py_DECREF(tmp);
}
}
Expand Down Expand Up @@ -2368,6 +2374,7 @@ strong_cache_free(StrongCacheNode *root)
static void
remove_from_strong_cache(zoneinfo_state *state, StrongCacheNode *node)
{
_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(state->ZoneInfoType);
if (state->ZONEINFO_STRONG_CACHE == node) {
state->ZONEINFO_STRONG_CACHE = node->next;
}
Expand Down Expand Up @@ -2422,6 +2429,7 @@ eject_from_strong_cache(zoneinfo_state *state, const PyTypeObject *const type,
return 0;
}

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(state->ZoneInfoType);
StrongCacheNode *cache = state->ZONEINFO_STRONG_CACHE;
StrongCacheNode *node = find_in_strong_cache(cache, key);
if (node != NULL) {
Expand Down Expand Up @@ -2478,6 +2486,7 @@ zone_from_strong_cache(zoneinfo_state *state, const PyTypeObject *const type,
return NULL; // Strong cache currently only implemented for base class
}

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(state->ZoneInfoType);
StrongCacheNode *cache = state->ZONEINFO_STRONG_CACHE;
StrongCacheNode *node = find_in_strong_cache(cache, key);

Expand All @@ -2504,6 +2513,7 @@ update_strong_cache(zoneinfo_state *state, const PyTypeObject *const type,
return;
}

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(state->ZoneInfoType);
StrongCacheNode *new_node = strong_cache_node_new(key, zone);
if (new_node == NULL) {
return;
Expand Down Expand Up @@ -2631,7 +2641,7 @@ static PyType_Slot zoneinfo_slots[] = {
{Py_tp_getattro, PyObject_GenericGetAttr},
{Py_tp_methods, zoneinfo_methods},
{Py_tp_members, zoneinfo_members},
{Py_tp_new, zoneinfo_new},
{Py_tp_new, zoneinfo_ZoneInfo},
{Py_tp_dealloc, zoneinfo_dealloc},
{Py_tp_traverse, zoneinfo_traverse},
{Py_tp_clear, zoneinfo_clear},
Expand Down
61 changes: 60 additions & 1 deletion Modules/clinic/_zoneinfo.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit bcaa724

Please sign in to comment.