Skip to content

Commit c51437d

Browse files
committed
Lock shared keys in Py_dict_lookup and use thread-safe lookup in insertdict
1 parent b4fe02f commit c51437d

File tree

1 file changed

+56
-2
lines changed

1 file changed

+56
-2
lines changed

Objects/dictobject.c

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,16 @@ ASSERT_DICT_LOCKED(PyObject *op)
162162
assert(_Py_IsOwnedByCurrentThread((PyObject *)mp) || IS_DICT_SHARED(mp));
163163
#define LOAD_KEYS_NENTRIES(d)
164164

165+
#define LOCK_KEYS_IF_SPLIT(keys, kind) \
166+
if (kind == DICT_KEYS_SPLIT) { \
167+
LOCK_KEYS(dk); \
168+
}
169+
170+
#define UNLOCK_KEYS_IF_SPLIT(keys, kind) \
171+
if (kind == DICT_KEYS_SPLIT) { \
172+
UNLOCK_KEYS(dk); \
173+
}
174+
165175
static inline Py_ssize_t
166176
load_keys_nentries(PyDictObject *mp)
167177
{
@@ -195,6 +205,9 @@ set_values(PyDictObject *mp, PyDictValues *values)
195205
#define DECREF_KEYS(dk) _Py_atomic_add_ssize(&dk->dk_refcnt, -1)
196206
#define LOAD_KEYS_NENTIRES(keys) _Py_atomic_load_ssize_relaxed(&keys->dk_nentries)
197207

208+
#define INCREF_KEYS_FT(dk) INCREF_KEYS(dk)
209+
#define DECREF_KEYS_FT(dk) INCREF_KEYS(dk)
210+
198211
static inline void split_keys_entry_added(PyDictKeysObject *keys)
199212
{
200213
ASSERT_KEYS_LOCKED(keys);
@@ -216,6 +229,8 @@ static inline void split_keys_entry_added(PyDictKeysObject *keys)
216229
#define INCREF_KEYS(dk) dk->dk_refcnt++
217230
#define DECREF_KEYS(dk) dk->dk_refcnt--
218231
#define LOAD_KEYS_NENTIRES(keys) keys->dk_nentries
232+
#define LOCK_KEYS_IF_SPLIT(keys, kind)
233+
#define UNLOCK_KEYS_IF_SPLIT(keys, kind)
219234
#define IS_DICT_SHARED(mp) (false)
220235
#define SET_DICT_SHARED(mp)
221236
#define LOAD_INDEX(keys, size, idx) ((const int##size##_t*)(keys->dk_indices))[idx]
@@ -1192,17 +1207,24 @@ _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **valu
11921207
PyDictKeysObject *dk;
11931208
DictKeysKind kind;
11941209
Py_ssize_t ix;
1195-
// TODO: Thread safety
11961210
start:
11971211
dk = mp->ma_keys;
11981212
kind = dk->dk_kind;
11991213

12001214
if (kind != DICT_KEYS_GENERAL) {
12011215
if (PyUnicode_CheckExact(key)) {
1216+
LOCK_KEYS_IF_SPLIT(dk, kind);
12021217
ix = unicodekeys_lookup_unicode(dk, key, hash);
1218+
UNLOCK_KEYS_IF_SPLIT(dk, kind);
12031219
}
12041220
else {
1221+
INCREF_KEYS_FT(dk);
1222+
LOCK_KEYS_IF_SPLIT(dk, kind);
1223+
12051224
ix = unicodekeys_lookup_generic(mp, dk, key, hash);
1225+
1226+
UNLOCK_KEYS_IF_SPLIT(dk, kind);
1227+
DECREF_KEYS_FT(dk);
12061228
if (ix == DKIX_KEY_CHANGED) {
12071229
goto start;
12081230
}
@@ -1443,6 +1465,7 @@ _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyOb
14431465
Py_DECREF(value);
14441466
goto read_failed;
14451467
}
1468+
14461469
}
14471470
}
14481471
else {
@@ -1709,7 +1732,38 @@ insertdict(PyInterpreterState *interp, PyDictObject *mp,
17091732
assert(mp->ma_keys->dk_kind == DICT_KEYS_GENERAL);
17101733
}
17111734

1712-
Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, &old_value);
1735+
Py_ssize_t ix;
1736+
#ifdef Py_GIL_DISABLED
1737+
PyDictKeysObject *dk = mp->ma_keys;
1738+
DictKeysKind kind = dk->dk_kind;
1739+
1740+
if (kind == DICT_KEYS_SPLIT) {
1741+
// Split keys can be mutated by other dictionaries, so we need
1742+
// to do a threadsafe lookup here. This is basically the split-key
1743+
// half of _Py_dict_lookup_threadsafe and avoids locking the
1744+
// shared keys if we can
1745+
if (PyUnicode_CheckExact(key)) {
1746+
ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash);
1747+
}
1748+
else {
1749+
ix = unicodekeys_lookup_generic_threadsafe(mp, dk, key, hash);
1750+
}
1751+
1752+
if (ix >= 0) {
1753+
old_value = mp->ma_values->values[ix];
1754+
}
1755+
else if (ix == DKIX_KEY_CHANGED) {
1756+
ix = _Py_dict_lookup(mp, key, hash, &old_value);
1757+
}
1758+
else {
1759+
old_value = NULL;
1760+
}
1761+
} else {
1762+
ix = _Py_dict_lookup(mp, key, hash, &old_value);
1763+
}
1764+
#else
1765+
ix = _Py_dict_lookup(mp, key, hash, &old_value);
1766+
#endif
17131767
if (ix == DKIX_ERROR)
17141768
goto Fail;
17151769

0 commit comments

Comments
 (0)