Added lock_shard guard to keep_alive (#737)

wjakob · Sep 28, 2024 · c1bab7e · c1bab7e
1 parent 7fad726
commit c1bab7e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/nb_type.cpp b/src/nb_type.cpp
@@ -1551,6 +1551,7 @@ void keep_alive(PyObject *nurse, PyObject *patient) {
     if (nb_type_check((PyObject *) Py_TYPE(nurse))) {
 #if defined(NB_FREE_THREADED)
         nb_shard &shard = internals->shard(inst_ptr((nb_inst *) nurse));
+        lock_shard guard(shard);
 #else
         nb_shard &shard = internals->shards[0];
 #endif
@@ -1604,6 +1605,7 @@ void keep_alive(PyObject *nurse, void *payload,
     if (nb_type_check((PyObject *) Py_TYPE(nurse))) {
 #if defined(NB_FREE_THREADED)
         nb_shard &shard = internals->shard(inst_ptr((nb_inst *) nurse));
+        lock_shard guard(shard);
 #else
         nb_shard &shard = internals->shards[0];
 #endif