Add locking for python free-threading

Python free-threading requires object locking around any mutable objects. Locks such as absl::Mutex may be used, or a python object's internal mutex may be used.  pybind11, unfortunately, does not have a built-in capability to ensure these locks are acquired.

This change adds the following types which aid in locking.
* ScopedPyCriticalSection and the ScopedPyCriticalSection2.
  These scoped locks acquire a python object's internal mutex.
* with_handle<T&> which provide access to the py::handle along with
  a type reference (which can be used to call ScopedPyCriticalSection)
* locking_type_caster, which acquires python locks when making a copy
  of an object into C++ (via T / T&& casts from the pybind11 type_caster).
* Additional concurrent tests.

The most common pattern for locked use is the following:
[](with_handle<T&> self, ...) {
  ScopedPyCriticalSection cs(self.handle.ptr());
  ... use self.value as needed ...
}

Additionally, most call sites have been audited to try and ensure correct
mutex use in the free threaded python environment. This has led to an
added template parameter, WithLocking, to several template functions.

Note that mutable types passed using by-value semantics are the likely to
require extra scrutiny, such as the use of locking_type_caster<>, or other
changes to ensure concurrent access is properly guarded.

See: #218
See: #254
PiperOrigin-RevId: 827716276
Change-Id: Id28d99b0a0ffb978988155892f8602c3c0d127f7

Author: laramiel

python/tensorstore/BUILD

@@ -228,6 +228,7 @@ pybind11_cc_library(
     srcs = ["json_type_caster.cc"],
     hdrs = ["json_type_caster.h"],
     deps = [
+        ":critical_section",
         "@com_github_pybind_pybind11//:pybind11",
         "@nlohmann_json//:json",
     ],
@@ -296,6 +297,7 @@ pybind11_cc_library(
     srcs = ["dim_expression.cc"],
     hdrs = ["dim_expression.h"],
     deps = [
+        ":critical_section",
         ":index",
         ":numpy_indexing_spec",
         ":sequence_parameter",
@@ -343,11 +345,13 @@ pybind11_cc_library(
     hdrs = ["index_space.h"],
     deps = [
         ":array_type_caster",
+        ":critical_section",
         ":dim_expression",
         ":gil_safe",
         ":homogeneous_tuple",
         ":index",
         ":json_type_caster",
+        ":locking_type_casters",
         ":numpy",
         ":numpy_indexing_spec",
         ":python_imports",
@@ -440,19 +444,21 @@ pybind11_cc_library(
     deps = [
         ":array_type_caster",
         ":context",
+        ":critical_section",
         ":data_type",
         ":define_heap_type",
         ":garbage_collection",
         ":homogeneous_tuple",
-        ":index",
         ":index_space",
         ":intrusive_ptr_holder",
         ":json_type_caster",
         ":keyword_arguments",
         ":kvstore",
+        ":locking_type_casters",
         ":result_type_caster",
         ":sequence_parameter",
         ":serialization",
+        ":status",
         ":tensorstore_module_components",
         ":unit",
         "//tensorstore:array",
@@ -462,7 +468,9 @@ pybind11_cc_library(
         "//tensorstore:data_type",
         "//tensorstore:index",
         "//tensorstore:json_serialization_options",
+        "//tensorstore:json_serialization_options_base",
         "//tensorstore:open_mode",
+        "//tensorstore:open_options",
         "//tensorstore:rank",
         "//tensorstore:schema",
         "//tensorstore:spec",
@@ -472,9 +480,10 @@ pybind11_cc_library(
         "//tensorstore/internal:global_initializer",
         "//tensorstore/internal:intrusive_ptr",
         "//tensorstore/internal/json:pprint_python",
+        "//tensorstore/internal/meta:type_traits",
+        "//tensorstore/kvstore",
         "//tensorstore/util:executor",
-        "//tensorstore/util:option",
-        "//tensorstore/util:result",
+        "//tensorstore/util:span",
         "//tensorstore/util:str_cat",
         "//tensorstore/util:unit",
         "@abseil-cpp//absl/status",
@@ -487,12 +496,8 @@ pybind11_cc_library(
 
 pybind11_cc_library(
     name = "data_type",
-    srcs = [
-        "data_type.cc",
-    ],
-    hdrs = [
-        "data_type.h",
-    ],
+    srcs = ["data_type.cc"],
+    hdrs = ["data_type.h"],
     deps = [
         ":json_type_caster",
         ":numpy",
@@ -520,6 +525,7 @@ pybind11_cc_library(
         ":array_type_caster",
         ":batch",
         ":context",
+        ":critical_section",
         ":data_type",
         ":define_heap_type",
         ":future",
@@ -531,6 +537,7 @@ pybind11_cc_library(
         ":json_type_caster",
         ":keyword_arguments",
         ":kvstore",
+        ":locking_type_casters",
         ":result_type_caster",
         ":sequence_parameter",
         ":serialization",
@@ -599,6 +606,7 @@ pybind11_cc_library(
         "//tensorstore/serialization",
         "//tensorstore/util:executor",
         "@com_github_pybind_pybind11//:pybind11",
+        "@nlohmann_json//:json",
     ],
     alwayslink = True,
 )
@@ -618,13 +626,38 @@ pybind11_cc_library(
     ],
 )
 
+pybind11_cc_library(
+    name = "critical_section",
+    hdrs = [
+        "critical_section.h",
+        "with_handle.h",
+    ],
+    deps = ["@com_github_pybind_pybind11//:pybind11"],
+)
+
+pybind11_cc_library(
+    name = "locking_type_casters",
+    hdrs = ["locking_type_casters.h"],
+    deps = [
+        ":critical_section",
+        "//tensorstore:array_storage_statistics",
+        "//tensorstore:batch",
+        "//tensorstore:chunk_layout",
+        "//tensorstore:index_interval",
+        "//tensorstore:schema",
+        "//tensorstore/kvstore:generation",
+        "//tensorstore/kvstore:key_range",
+        "@com_github_pybind_pybind11//:pybind11",
+    ],
+)
+
 pybind11_cc_library(
     name = "garbage_collection",
     srcs = ["garbage_collection.cc"],
     hdrs = ["garbage_collection.h"],
     deps = [
+        ":critical_section",
         ":define_heap_type",
-        ":gil_safe",
         ":tensorstore_module_components",
         "//tensorstore/internal:global_initializer",
         "//tensorstore/internal:intrusive_ptr",
@@ -642,6 +675,7 @@ pybind11_cc_library(
     srcs = ["future.cc"],
     hdrs = ["future.h"],
     deps = [
+        ":critical_section",
         ":define_heap_type",
         ":garbage_collection",
         ":gil_safe",
@@ -658,8 +692,10 @@ pybind11_cc_library(
         "//tensorstore/util:executor",
         "//tensorstore/util:future",
         "//tensorstore/util:result",
+        "@abseil-cpp//absl/base:core_headers",
         "@abseil-cpp//absl/functional:function_ref",
         "@abseil-cpp//absl/status",
+        "@abseil-cpp//absl/synchronization",
         "@abseil-cpp//absl/time",
         "@com_github_pybind_pybind11//:pybind11",
     ],
@@ -710,6 +746,7 @@ pybind11_cc_library(
     hdrs = ["numpy_indexing_spec.h"],
     deps = [
         ":array_type_caster",
+        ":critical_section",
         ":data_type",
         ":index",
         ":numpy",
@@ -723,15 +760,13 @@ pybind11_cc_library(
         "//tensorstore:index",
         "//tensorstore:rank",
         "//tensorstore:static_cast",
-        "//tensorstore/index_space:dimension_identifier",
-        "//tensorstore/index_space:dimension_index_buffer",
-        "//tensorstore/index_space:index_transform",
         "//tensorstore/index_space:numpy_indexing_spec",
         "//tensorstore/util:byte_strided_pointer",
         "//tensorstore/util:iterate",
         "//tensorstore/util:span",
         "//tensorstore/util:str_cat",
         "@abseil-cpp//absl/base:core_headers",
+        "@abseil-cpp//absl/meta:type_traits",
         "@abseil-cpp//absl/status",
         "@com_github_pybind_pybind11//:pybind11",
     ],
@@ -808,7 +843,7 @@ tensorstore_pytest_test(
     name = "future_test",
     size = "medium",
     srcs = ["tests/future_test.py"],
-    tags = ["cpu:2"],
+    tags = ["cpu:4"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -819,6 +854,7 @@ tensorstore_pytest_test(
     name = "chunk_layout_test",
     size = "small",
     srcs = ["tests/chunk_layout_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -870,6 +906,8 @@ pybind11_cc_library(
     srcs = ["batch.cc"],
     hdrs = ["batch.h"],
     deps = [
+        ":critical_section",
+        ":locking_type_casters",
         ":tensorstore_module_components",
         "//tensorstore:batch",
         "//tensorstore/internal:global_initializer",
@@ -955,11 +993,13 @@ pybind11_cc_library(
     hdrs = ["chunk_layout.h"],
     deps = [
         ":array_type_caster",
+        ":critical_section",
         ":homogeneous_tuple",
         ":index",
         ":index_space",
         ":json_type_caster",
         ":keyword_arguments",
+        ":locking_type_casters",
         ":result_type_caster",
         ":sequence_parameter",
         ":serialization",
@@ -979,6 +1019,7 @@ pybind11_cc_library(
         "//tensorstore/util:str_cat",
         "@abseil-cpp//absl/status",
         "@com_github_pybind_pybind11//:pybind11",
+        "@nlohmann_json//:json",
     ],
     alwayslink = True,
 )
@@ -1041,6 +1082,7 @@ pybind11_cc_library(
     srcs = ["serialization.cc"],
     hdrs = ["serialization.h"],
     deps = [
+        ":critical_section",
         ":garbage_collection",
         ":gil_safe",
         ":result_type_caster",
@@ -1059,6 +1101,7 @@ pybind11_cc_library(
         "@abseil-cpp//absl/functional:function_ref",
         "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/strings:cord",
+        "@abseil-cpp//absl/synchronization",
         "@com_github_pybind_pybind11//:pybind11",
         "@riegeli//riegeli/bytes:cord_writer",
         "@riegeli//riegeli/bytes:string_reader",
@@ -1075,11 +1118,13 @@ pybind11_cc_library(
     deps = [
         ":batch",
         ":context",
+        ":critical_section",
         ":define_heap_type",
         ":future",
         ":garbage_collection",
         ":json_type_caster",
         ":keyword_arguments",
+        ":locking_type_casters",
         ":result_type_caster",
         ":serialization",
         ":status",
@@ -1100,6 +1145,7 @@ pybind11_cc_library(
         "//tensorstore/kvstore:key_range",
         "//tensorstore/util:executor",
         "//tensorstore/util:future",
+        "//tensorstore/util:result",
         "//tensorstore/util:str_cat",
         "//tensorstore/util/garbage_collection",
         "@abseil-cpp//absl/status",
@@ -1165,6 +1211,7 @@ pybind11_cc_library(
         ":garbage_collection",
         ":gil_safe",
         ":keyword_arguments",
+        ":locking_type_casters",
         ":result_type_caster",
         ":serialization",
         ":spec",

python/tensorstore/batch.cc

@@ -22,7 +22,9 @@
 #include <optional>
 #include <utility>
 
+#include "python/tensorstore/critical_section.h"
 #include "python/tensorstore/tensorstore_module_components.h"
+#include "python/tensorstore/with_handle.h"
 #include "tensorstore/batch.h"
 #include "tensorstore/internal/global_initializer.h"
 #include "tensorstore/util/executor.h"
@@ -135,12 +137,15 @@ Operations
 }
 
 void DefineBatchAttributes(BatchCls& cls) {
-  using Self = Batch;
   cls.def(py::init([]() { return Batch::New(); }), R"(
 Creates a new batch.
 )");
   cls.def(
-      "submit", [](Self& self) { self.Release(); },
+      "submit",
+      [](with_handle<Batch&> self) {
+        ScopedPyCriticalSection cs(self.handle.ptr());
+        self.value.Release();
+      },
       R"(
 Submits the batch.
 
@@ -153,9 +158,15 @@ will result in an error.
 Group:
   Operations
 )");
-  cls.def("__enter__", [](const Self& self) { return &self; });
-  cls.def("__exit__", [](Self& self, py::object exc_type, py::object exc_value,
-                         py::object traceback) { self.Release(); });
+  cls.def("__enter__", [](with_handle<const Batch&> self) {
+    ScopedPyCriticalSection cs(self.handle.ptr());
+    return &self.value;
+  });
+  cls.def("__exit__", [](with_handle<Batch&> self, py::object exc_type,
+                         py::object exc_value, py::object traceback) {
+    ScopedPyCriticalSection cs(self.handle.ptr());
+    self.value.Release();
+  });
 }
 
 void RegisterBatchBindings(pybind11::module m, Executor defer) {

python/tensorstore/batch.h

@@ -21,6 +21,7 @@
 
 #include <optional>
 
+#include "python/tensorstore/locking_type_casters.h"  // IWYU pragma: keep
 #include "tensorstore/batch.h"
 
 namespace tensorstore {