Repository: man-group/ArcticDB (language: python)

Issue:
write batch produces segmentation fault when using string indexes and dynamic strings is enabled
### Describe the bug

When writing in batches and using string indexes with dynamic strings set to true, we call the function PyUnicode_AsUTF8String. However, this function is not thread-safe, and accessing it concurrently can cause segmentation faults. 

To prevent concurrent access to this function, we need to use mechanisms like the Python GIL locks.

### Steps/Code to Reproduce

```
def test_batch_write_segfault(s3_version_store_v1):
     lib = s3_version_store_v1
     syms = ["opt_hard_upper_posbounds", "estimated_wave_volume_usd"]

     index1 = [f"AST{i}" for i in range(586)]
     index2 = [f"AST{i}" for i in range(585)]

     vals1 = np.random.randn(586).astype('float64')
     vals2 = np.random.randn(585).astype('float64')

     data = [pd.Series(data=vals1, index=index1), pd.Series(data=vals2, index=index2)]

     # sometimes segfaults the first time, sometimes doesn't, so run in a loop and it is very likely to crash
     for i in range(5):
         lib.batch_write(syms, data)
```

### Expected Results

We do not expect segmentation faults.

### OS, Python Version and ArcticDB Version

Python: 3.6.8 (default, Jun 20 2023, 11:53:23) 
[GCC 4.8.5 20150623 (Red Hat 4.8.5-44)]
OS: Linux-5.15.0-76-generic-x86_64-with-centos-7.9.2009-Core

### Backend storage used

_No response_

### Additional Context

_No response_


Discussion / hints:
A prototype [solution](https://github.com/man-group/ArcticDB/pull/768) has been developed to fix this issue, but it is just a demonstration that concurrent access to PyUnicode_AsUTF8String is the problem.

Fix patch:
diff --git a/cpp/arcticdb/pipeline/frame_utils.hpp b/cpp/arcticdb/pipeline/frame_utils.hpp
index 0ae859b194..98b32f9278 100644
--- a/cpp/arcticdb/pipeline/frame_utils.hpp
+++ b/cpp/arcticdb/pipeline/frame_utils.hpp
@@ -15,6 +15,7 @@
 #include <arcticdb/pipeline/input_tensor_frame.hpp>
 #include <arcticdb/stream/protobuf_mappings.hpp>
 #include <arcticdb/entity/protobuf_mappings.hpp>
+#include <arcticdb/python/gil_lock.hpp>
 #include <arcticdb/python/python_types.hpp>
 #include <arcticdb/python/python_to_tensor_frame.hpp>
 #include <arcticdb/pipeline/string_pool_utils.hpp>
@@ -130,6 +131,11 @@ void aggregator_set_data(
                     ptr_data = flatten_tensor<PyObject*>(flattened_buffer, rows_to_write, tensor, slice_num, regular_slice_size);
 
                 auto none = py::none{};
+                // GIL will be acquired if there is a string that is not pure ASCII/UTF-8
+                // In this case a PyObject will be allocated by convert::py_unicode_to_buffer
+                // If such a string is encountered in a column, then the GIL will be held until that whole column has
+                // been processed, on the assumption that if a column has one such string it will probably have many.
+                std::optional<ScopedGILLock> scoped_gil_lock;
                 for (size_t s = 0; s < rows_to_write; ++s, ++ptr_data) {
                     if (*ptr_data == none.ptr()) {
                         agg.set_no_string_at(col, s, not_a_string());
@@ -137,7 +143,7 @@ void aggregator_set_data(
                         agg.set_no_string_at(col, s, nan_placeholder());
                     } else {
                         if constexpr (is_utf_type(slice_value_type(dt))) {
-                            auto wrapper= convert::py_unicode_to_buffer(*ptr_data);
+                            auto wrapper= convert::py_unicode_to_buffer(*ptr_data, scoped_gil_lock);
                             agg.set_string_at(col, s, wrapper.buffer_, wrapper.length_);
                         } else {
                             auto wrapper = convert::pystring_to_buffer(*ptr_data);
diff --git a/cpp/arcticdb/python/python_to_tensor_frame.cpp b/cpp/arcticdb/python/python_to_tensor_frame.cpp
index bb885a924d..8c0f41c387 100644
--- a/cpp/arcticdb/python/python_to_tensor_frame.cpp
+++ b/cpp/arcticdb/python/python_to_tensor_frame.cpp
@@ -24,22 +24,34 @@ bool is_unicode(PyObject *obj) {
     return PyUnicode_Check(obj);
 }
 
-PyStringWrapper pystring_to_buffer(PyObject *obj, py::handle handle) {
+PyStringWrapper pystring_to_buffer(PyObject *obj) {
     char *buffer;
     ssize_t length;
     util::check(!is_unicode(obj), "Unexpected unicode object");
     if (PYBIND11_BYTES_AS_STRING_AND_SIZE(obj, &buffer, &length))
         util::raise_rte("Unable to extract string contents! (invalid type)");
 
-    return {buffer, length, handle};
+    return {buffer, length, obj};
 }
 
-PyStringWrapper py_unicode_to_buffer(PyObject *obj) {
+PyStringWrapper py_unicode_to_buffer(PyObject *obj, std::optional<ScopedGILLock>& scoped_gil_lock) {
     if (is_unicode(obj)) {
-        py::handle handle{PyUnicode_AsUTF8String(obj)};
-        if (!handle)
-            util::raise_rte("Unable to extract string contents! (encoding issue)");
-        return pystring_to_buffer(handle.ptr(), handle);
+        if (PyUnicode_IS_COMPACT_ASCII(obj)) {
+            return {reinterpret_cast<char *>(PyUnicode_DATA(obj)), PyUnicode_GET_LENGTH(obj)};
+        // Later versions of cpython expose macros in unicodeobject.h to perform this check, and to get the utf8_length,
+        // but for 3.6 we have to hand-roll it
+        } else if (reinterpret_cast<PyCompactUnicodeObject*>(obj)->utf8) {
+            return {reinterpret_cast<PyCompactUnicodeObject*>(obj)->utf8, reinterpret_cast<PyCompactUnicodeObject*>(obj)->utf8_length};
+        } else {
+            internal::check<ErrorCode::E_ASSERTION_FAILURE>(PyUnicode_READY(obj) == 0, "PyUnicode_READY failed");
+            if (!scoped_gil_lock.has_value()) {
+                scoped_gil_lock.emplace();
+            }
+            PyObject* utf8_obj = PyUnicode_AsUTF8String(obj);
+            if (!utf8_obj)
+                util::raise_rte("Unable to extract string contents! (encoding issue)");
+            return pystring_to_buffer(utf8_obj);
+        }
     } else {
         util::raise_rte("Expected unicode");
     }
diff --git a/cpp/arcticdb/python/python_to_tensor_frame.hpp b/cpp/arcticdb/python/python_to_tensor_frame.hpp
index ebb1e61441..b08012d569 100644
--- a/cpp/arcticdb/python/python_to_tensor_frame.hpp
+++ b/cpp/arcticdb/python/python_to_tensor_frame.hpp
@@ -8,6 +8,7 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include <arcticdb/python/gil_lock.hpp>
 #include <arcticdb/pipeline/input_tensor_frame.hpp>
 #include <arcticdb/entity/native_tensor.hpp>
 #include <string>
@@ -20,25 +21,28 @@ using namespace arcticdb::entity;
 struct ARCTICDB_VISIBILITY_HIDDEN PyStringWrapper {
     char *buffer_;
     size_t length_;
-    py::handle handle_;
+    PyObject* obj_;
 
-    PyStringWrapper(char *buf, ssize_t len, py::handle handle) :
+    // If the underlying Python string is ASCII or UTF-8, we can use the input object's underlying buffer, and obj will
+    // be nullptr in this ctor. For unicode, the Python C API method is used to construct a new Python object which must
+    // be DECREFFed on destruction to free the underlying memory.
+    PyStringWrapper(char *buf, ssize_t len, PyObject* obj=nullptr) :
         buffer_(buf),
         length_(size_t(len)),
-        handle_(handle) {}
+        obj_(obj) {}
 
     ~PyStringWrapper() {
-        if (handle_)
-            handle_.dec_ref();
+        if (obj_)
+            Py_DECREF(obj_);
     }
 };
 
 PyStringWrapper pystring_to_buffer(
-    PyObject *obj,
-    py::handle handle = py::handle());
+    PyObject *obj);
 
 PyStringWrapper py_unicode_to_buffer(
-    PyObject *obj);
+    PyObject *obj,
+    std::optional<ScopedGILLock>& scoped_gil_lock);
 
 NativeTensor obj_to_tensor(PyObject *ptr);
 
diff --git a/cpp/arcticdb/version/local_versioned_engine.cpp b/cpp/arcticdb/version/local_versioned_engine.cpp
index 22dddc13cb..30793cd4f3 100644
--- a/cpp/arcticdb/version/local_versioned_engine.cpp
+++ b/cpp/arcticdb/version/local_versioned_engine.cpp
@@ -1337,6 +1337,8 @@ std::vector<std::variant<VersionedItem, DataError>> LocalVersionedEngine::batch_
     bool validate_index,
     bool throw_on_error
 ) {
+    py::gil_scoped_release release_gil;
+
     auto write_options = get_write_options();
     auto update_info_futs = batch_get_latest_undeleted_version_and_next_version_id_async(store(),
                                                                                          version_map(),
@@ -1447,6 +1449,7 @@ std::vector<std::variant<VersionedItem, DataError>> LocalVersionedEngine::batch_
     bool validate_index,
     bool upsert,
     bool throw_on_error) {
+    py::gil_scoped_release release_gil;
 
     auto stream_update_info_futures = batch_get_latest_undeleted_version_and_next_version_id_async(store(),
                                                                                                     version_map(),


Test patch:
diff --git a/cpp/arcticdb/version/test/test_sparse.cpp b/cpp/arcticdb/version/test/test_sparse.cpp
index bd3814bf87..b252dc1522 100644
--- a/cpp/arcticdb/version/test/test_sparse.cpp
+++ b/cpp/arcticdb/version/test/test_sparse.cpp
@@ -182,7 +182,8 @@ TEST_F(SparseTestStore, SimpleRoundtripStrings) {
 
     ASSERT_EQ(frame.row_count(), 2);
     auto val1 = frame.scalar_at<PyObject*>(0, 1);
-    auto str_wrapper = convert::py_unicode_to_buffer(val1.value());
+    std::optional<ScopedGILLock> scoped_gil_lock;
+    auto str_wrapper = convert::py_unicode_to_buffer(val1.value(), scoped_gil_lock);
     ASSERT_EQ(strcmp(str_wrapper.buffer_, "five"), 0);
 
     auto val2 = frame.string_at(0, 2);
@@ -560,7 +561,8 @@ TEST_F(SparseTestStore, CompactWithStrings) {
 
     for(size_t i = 0; i < num_rows; i += 2) {
         auto val1 = frame.scalar_at<PyObject*>(i, 1);
-        auto str_wrapper = convert::py_unicode_to_buffer(val1.value());
+        std::optional<ScopedGILLock> scoped_gil_lock;
+        auto str_wrapper = convert::py_unicode_to_buffer(val1.value(), scoped_gil_lock);
         auto expected = fmt::format("{}", i + 1);
         ASSERT_EQ(strcmp(str_wrapper.buffer_, expected.data()), 0);
 
diff --git a/python/tests/nonreg/arcticdb/version_store/test_nonreg_specific.py b/python/tests/nonreg/arcticdb/version_store/test_nonreg_specific.py
index ed964416d5..d3e7ae09aa 100644
--- a/python/tests/nonreg/arcticdb/version_store/test_nonreg_specific.py
+++ b/python/tests/nonreg/arcticdb/version_store/test_nonreg_specific.py
@@ -241,3 +241,31 @@ def test_update_fixed_width_to_dynamic_strings(lib_type, request):
     expected_df = dynamic_strings_data
     read_df = lib.read(symbol).data
     assert_frame_equal(expected_df, read_df)
+
+
+# https://github.com/man-group/ArcticDB/issues/767
+# Batch write and append call aggregator_set_data from threads running in the CPU threadpool (i.e. not the main thread)
+# With unicode strings, a PyObject allocation is needed, and therefore the thread must be holding the GIL
+# This test ensures that the correct thread is holding the GIL when performing these allocations, and that there is no
+# deadlock.
+# This is not an issue with non-batch methods as aggregator_set_data is called from the main thread, which pybind11
+# ensures is holding the GIL on entry to the C++ layer.
+def test_batch_write_unicode_strings(lmdb_version_store):
+    lib = lmdb_version_store
+    syms = ["sym1", "sym2"]
+    # 10 was too small to trigger problem
+    num_rows = 100
+
+    index = np.arange(num_rows)
+    u_umlaut = b"\xc3\x9c".decode("utf-8")
+    unicode_vals = [u_umlaut] * num_rows
+
+    data = [
+        pd.Series(data=unicode_vals, index=index),
+        pd.Series(data=unicode_vals, index=index),
+    ]
+
+    # The problem was not always triggered on the first call
+    for _ in range(5):
+        lib.batch_write(syms, data)
+        lib.batch_append(syms, data)