Fix a data race between clear_keep_alive and state. (#1191)

In JAX CI, we saw the following data race:

```
WARNING: ThreadSanitizer: data race (pid=36276)
  Read of size 4 at 0x7ad434c76774 by thread T4 (mutexes: read M0):
    #0 nanobind::detail::nb_type_get(std::type_info const*, _object*, unsigned char, nanobind::detail::cleanup_list*, void**) /proc/self/cwd/external/nanobind/src/nb_type.cpp:1536:17 (libjax_common.so+0xa5991a) (BuildId: 2414288b76007439f199e80c1b7b912642188753)
    #1 nanobind::detail::type_caster_base::from_python(nanobind::handle, unsigned char, nanobind::detail::cleanup_list*) /proc/self/cwd/external/nanobind/include/nanobind/nb_cast.h:481:16 (libjax_common.so+0x5d96184) (BuildId: 2414288b76007439f199e80c1b7b912642188753)
    #2 _object* nanobind::detail::func_create(xla::Layout const& (xla::PjRtLayout::*)() const, nanobind::scope const&, nanobind::name const&, nanobind::is_method const&)::'lambda'(xla::PjRtLayout const*), xla::Layout const&, xla::PjRtLayout const*, 0ul, nanobind::scope, nanobind::name, nanobind::is_method>(xla::PjRtLayout&&, xla::Layout const& (*)(nanobind::scope, nanobind::name, nanobind::is_method), std::integer_sequence, nanobind::scope const&, nanobind::name const&, nanobind::is_method const&)::'lambda'(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*)::operator()(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*) const /proc/self/cwd/external/nanobind/include/nanobind/nb_func.h:254:41 (libjax_common.so+0x5d96184)

...

  Previous write of size 4 at 0x7ad434c76774 by thread T5 (mutexes: read M0):
    #0 nanobind::detail::keep_alive(_object*, void*, void (*)(void*) noexcept) /proc/self/cwd/external/nanobind/src/nb_type.cpp:1661:47 (libjax_common.so+0xa5a508) (BuildId: 2414288b76007439f199e80c1b7b912642188753)
    #1 nanobind::detail::shared_from_cpp(std::shared_ptr&&, _object*) /proc/self/cwd/external/nanobind/include/nanobind/stl/shared_ptr.h:56:5 (libjax_common.so+0x8c6bc6) (BuildId: 2414288b76007439f199e80c1b7b912642188753)
    #2 nanobind::detail::type_caster, int>::from_cpp(std::shared_ptr const&, nanobind::rv_policy, nanobind::detail::cleanup_list*) /proc/self/cwd/external/nanobind/include/nanobind/stl/shared_ptr.h:129:13 (libjax_common.so+0x8c6899) (BuildId: 2414288b76007439f199e80c1b7b912642188753)
    #3 _object* nanobind::detail::func_create> (), jax::PyArray>, std::shared_ptr, jax::PyArray&, 0ul, nanobind::is_method, nanobind::is_getter>(xla::ValueOrThrowWrapper> (), jax::PyArray>&&, std::shared_ptr (*)(jax::PyArray&), std::integer_sequence, nanobind::is_method const&, nanobind::is_getter const&)::'lambda'(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*)::operator()(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*) const /proc/self/cwd/external/nanobind/include/nanobind/nb_func.h:274:22 (libjax_common.so+0x8c63d5) (BuildId: 2414288b76007439f199e80c1b7b912642188753)

```

The race looks like this:
* two threads both attempt to return the same std::shared_ptr<> object
  concurrently.
* thread (a) calls nb_type_put and gets the is_new=true case. Thread (b)
  calls nb_type_put and gets the is_new=false case.
* thread (b) returns the object to Python, and then passes that object
  as an argument to a nanobind-bound function, which reads inst->state.
* thread (a) reaches the is_new branch and calls keep_alive, which
  writes inst->clear_keep_alive.

`state` and `clear_keep_alive` are two bitfields on the same object
accessed concurrently, so this is a data race, per the C++20 standard:

[intro.races] (6.9.2.1) Data races: Paragraph 21 defines a data race and its consequence: "The execution of a program contains a data race if it contains two potentially concurrent conflicting actions, at least one of which is not atomic, and neither happens before the other... Any such data race results in undefined behavior.". Paragraph 2 defines conflicting actions as modifying or reading/modifying the same memory location.

[intro.memory] (6.7.1) Memory model: This section precisely defines what a "memory location" is for bitfields. Paragraph 3 states: "A memory location is either an object of scalar type or a maximal sequence of adjacent bit-fields all having nonzero width.".

This PR splits clear_keep_alive into a separate non-bitfield member.
This may or may not be a sufficient fix, in particular, one can imagine
that it may be possible for multiple threads to call keep_alive
concurrently, so it may need to be an atomic value.

Author: hawkinsp

docs/changelog.rst

@@ -18,6 +18,10 @@ below inherit that of the preceding release.
 Version TBD (not yet released)
 ------------------------------
 
+- ABI version 17. The layout of the `nb_inst` class was changed to fix a data race
+  between `clear_keep_alive` and other bitfields in the same struct (PR `#1191
+  <https://github.com/wjakob/nanobind/pull/1191>`__)
+
 - Nanobind now officially supports **MinGW-w64** and **Intel ICX** (the modern
   Clang-based Intel compiler). Continuous integration tests have been added to
   ensure compatibility with these compilers on an ongoing basis.

src/nb_abi.h

@@ -14,7 +14,7 @@
 
 /// Tracks the version of nanobind's internal data structures
 #ifndef NB_INTERNALS_VERSION
-#  define NB_INTERNALS_VERSION 16
+#  define NB_INTERNALS_VERSION 17
 #endif
 
 #if defined(__MINGW32__)

src/nb_internals.h

@@ -56,7 +56,7 @@ struct nb_inst { // usually: 24 bytes
 
     /// State of the C++ object this instance points to: is it constructed?
     /// can we use it?
-    uint32_t state : 2;
+    uint8_t state : 2;
 
     // Values for `state`. Note that the numeric values of these are relied upon
     // for an optimization in `nb_type_get()`.
@@ -70,25 +70,27 @@ struct nb_inst { // usually: 24 bytes
      * relative offset to a pointer that must be dereferenced to get to the
      * instance data. 'direct' is 'true' in the former case.
      */
-    uint32_t direct : 1;
+    uint8_t direct : 1;
 
     /// Is the instance data co-located with the Python object?
-    uint32_t internal : 1;
+    uint8_t internal : 1;
 
     /// Should the destructor be called when this instance is GCed?
-    uint32_t destruct : 1;
+    uint8_t destruct : 1;
 
     /// Should nanobind call 'operator delete' when this instance is GCed?
-    uint32_t cpp_delete : 1;
-
-    /// Does this instance hold references to others? (via internals.keep_alive)
-    uint32_t clear_keep_alive : 1;
+    uint8_t cpp_delete : 1;
 
     /// Does this instance use intrusive reference counting?
-    uint32_t intrusive : 1;
+    uint8_t intrusive : 1;
+
+    /// Does this instance hold references to others? (via internals.keep_alive)
+    /// This may be accessed concurrently to 'state', so it must not be in
+    /// the same bitfield as 'state'.
+    uint8_t clear_keep_alive;
 
     // That's a lot of unused space. I wonder if there is a good use for it..
-    uint32_t unused : 24;
+    uint16_t unused;
 };
 
 static_assert(sizeof(nb_inst) == sizeof(PyObject) + sizeof(uint32_t) * 2);

tests/test_thread.cpp

@@ -1,4 +1,8 @@
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/shared_ptr.h>
+
+#include <memory>
+#include <vector>
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -32,6 +36,11 @@ class ClassWithClassProperty {
     ClassWithProperty value_;
 };
 
+struct AnInt {
+    int value;
+    AnInt(int v) : value(v) {}
+};
+
 
 NB_MODULE(test_thread_ext, m) {
     nb::class_<Counter>(m, "Counter")
@@ -68,4 +77,17 @@ NB_MODULE(test_thread_ext, m) {
             new (self) ClassWithClassProperty(std::move(value));
           }, nb::arg("value"))
         .def_prop_ro("prop1", &ClassWithClassProperty::get_prop);
+
+    nb::class_<AnInt>(m, "AnInt")
+        .def(nb::init<int>())
+        .def_rw("value", &AnInt::value);
+
+    std::vector<std::shared_ptr<AnInt>> shared_ints;
+    for (int i = 0; i < 5; ++i) {
+        shared_ints.push_back(std::make_shared<AnInt>(i));
+    }
+    m.def("fetch_shared_int", [shared_ints](int i) {
+        return shared_ints.at(i);
+    });
+    m.def("consume_an_int", [](AnInt* p) { return p->value; });
 }

tests/test_thread.py

@@ -1,3 +1,6 @@
+import random
+import threading
+
 import test_thread_ext as t
 from test_thread_ext import Counter, GlobalData, ClassWithProperty, ClassWithClassProperty
 from common import parallelize
@@ -100,3 +103,16 @@ def f():
             _ = c2.prop1.prop2
 
     parallelize(f, n_threads=n_threads)
+
+
+def test08_shared_ptr_threaded_access(n_threads=8):
+    # Test for keep_alive racing with other fields.
+    def f(barrier):
+        i = random.randint(0, 4)
+        barrier.wait()
+        p = t.fetch_shared_int(i)
+        assert t.consume_an_int(p) == i
+
+    for _ in range(100):
+        barrier = threading.Barrier(n_threads)
+        parallelize(lambda: f(barrier), n_threads=n_threads)