You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by pa...@apache.org on 2024/02/16 21:19:51 UTC
(arrow-nanoarrow) branch main updated: fix: Ensure that the deallocator called by ArrowBufferDeallocator() is called exactly once (#387)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new dfa48d14 fix: Ensure that the deallocator called by ArrowBufferDeallocator() is called exactly once (#387)
dfa48d14 is described below
commit dfa48d1464e8ea2b34c6266f852313699def5175
Author: Dewey Dunnington <de...@dunnington.ca>
AuthorDate: Fri Feb 16 17:19:46 2024 -0400
fix: Ensure that the deallocator called by ArrowBufferDeallocator() is called exactly once (#387)
This PR makes it more predictable to work with custom allocators,
including the `ArrowBufferDeallocator` that powers the R and (soon)
Python "build from buffers" behaviour. It also improves testing for
using custom allocators and adds a templated C++ wrapper.
---
r/.Rbuildignore | 1 +
r/.gitignore | 1 +
src/nanoarrow/buffer_inline.h | 15 +++---
src/nanoarrow/nanoarrow.hpp | 52 ++++++++++++++++++++
src/nanoarrow/nanoarrow_hpp_test.cc | 64 +++++++++++++++++++++++++
src/nanoarrow/utils.c | 29 ++++++++----
src/nanoarrow/utils_test.cc | 94 ++++++++++++++++++++++++++++---------
7 files changed, 216 insertions(+), 40 deletions(-)
diff --git a/r/.Rbuildignore b/r/.Rbuildignore
index 6c40992c..8cbf5e33 100644
--- a/r/.Rbuildignore
+++ b/r/.Rbuildignore
@@ -11,3 +11,4 @@
^cran-comments\.md$
^bootstrap\.R$
^\.cache$
+^compile_commands\.json$
diff --git a/r/.gitignore b/r/.gitignore
index 0d397436..f629b0db 100644
--- a/r/.gitignore
+++ b/r/.gitignore
@@ -18,3 +18,4 @@
.Rproj.user
.Rhistory
docs
+compile_commands.json
diff --git a/src/nanoarrow/buffer_inline.h b/src/nanoarrow/buffer_inline.h
index e27ff67c..b5fc627b 100644
--- a/src/nanoarrow/buffer_inline.h
+++ b/src/nanoarrow/buffer_inline.h
@@ -46,6 +46,8 @@ static inline void ArrowBufferInit(struct ArrowBuffer* buffer) {
static inline ArrowErrorCode ArrowBufferSetAllocator(
struct ArrowBuffer* buffer, struct ArrowBufferAllocator allocator) {
+ // This is not a perfect test for "has a buffer already been allocated"
+ // but is likely to catch most cases.
if (buffer->data == NULL) {
buffer->allocator = allocator;
return NANOARROW_OK;
@@ -55,20 +57,15 @@ static inline ArrowErrorCode ArrowBufferSetAllocator(
}
static inline void ArrowBufferReset(struct ArrowBuffer* buffer) {
- if (buffer->data != NULL) {
- buffer->allocator.free(&buffer->allocator, (uint8_t*)buffer->data,
- buffer->capacity_bytes);
- buffer->data = NULL;
- }
-
- buffer->capacity_bytes = 0;
- buffer->size_bytes = 0;
+ buffer->allocator.free(&buffer->allocator, (uint8_t*)buffer->data,
+ buffer->capacity_bytes);
+ ArrowBufferInit(buffer);
}
static inline void ArrowBufferMove(struct ArrowBuffer* src, struct ArrowBuffer* dst) {
memcpy(dst, src, sizeof(struct ArrowBuffer));
src->data = NULL;
- ArrowBufferReset(src);
+ ArrowBufferInit(src);
}
static inline ArrowErrorCode ArrowBufferResize(struct ArrowBuffer* buffer,
diff --git a/src/nanoarrow/nanoarrow.hpp b/src/nanoarrow/nanoarrow.hpp
index 8d5b841e..09a03151 100644
--- a/src/nanoarrow/nanoarrow.hpp
+++ b/src/nanoarrow/nanoarrow.hpp
@@ -241,6 +241,13 @@ class Unique {
T data_;
};
+template <typename T>
+static inline void DeallocateWrappedBuffer(struct ArrowBufferAllocator* allocator,
+ uint8_t* ptr, int64_t size) {
+ auto obj = reinterpret_cast<T*>(allocator->private_data);
+ delete obj;
+}
+
/// @}
} // namespace internal
@@ -273,6 +280,51 @@ using UniqueArrayView = internal::Unique<struct ArrowArrayView>;
/// @}
+/// \defgroup nanoarrow_hpp-buffer Buffer helpers
+///
+/// Helpers to wrap buffer-like C++ objects as ArrowBuffer objects that can
+/// be used to build ArrowArray objects.
+///
+/// @{
+
+/// \brief Initialize a buffer wrapping an arbitrary C++ object
+///
+/// Initializes a buffer with a release callback that deletes the moved obj
+/// when ArrowBufferReset is called. This version is useful for wrapping
+/// an object whose .data() member is missing or unrelated to the buffer
+/// value that is destined for a the buffer of an ArrowArray. T must be movable.
+template <typename T>
+static inline void BufferInitWrapped(struct ArrowBuffer* buffer, T obj,
+ const uint8_t* data, int64_t size_bytes) {
+ T* obj_moved = new T(std::move(obj));
+ buffer->data = const_cast<uint8_t*>(data);
+ buffer->size_bytes = size_bytes;
+ buffer->capacity_bytes = 0;
+ buffer->allocator =
+ ArrowBufferDeallocator(&internal::DeallocateWrappedBuffer<T>, obj_moved);
+}
+
+/// \brief Initialize a buffer wrapping a C++ sequence
+///
+/// Specifically, this uses obj.data() to set the buffer address and
+/// obj.size() * sizeof(T::value_type) to set the buffer size. This works
+/// for STL containers like std::vector, std::array, and std::string.
+/// This function moves obj and ensures it is deleted when ArrowBufferReset
+/// is called.
+template <typename T>
+void BufferInitSequence(struct ArrowBuffer* buffer, T obj) {
+ // Move before calling .data() (matters sometimes).
+ T* obj_moved = new T(std::move(obj));
+ buffer->data =
+ const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(obj_moved->data()));
+ buffer->size_bytes = obj_moved->size() * sizeof(typename T::value_type);
+ buffer->capacity_bytes = 0;
+ buffer->allocator =
+ ArrowBufferDeallocator(&internal::DeallocateWrappedBuffer<T>, obj_moved);
+}
+
+/// @}
+
/// \defgroup nanoarrow_hpp-array-stream ArrayStream helpers
///
/// These classes provide simple ArrowArrayStream implementations that
diff --git a/src/nanoarrow/nanoarrow_hpp_test.cc b/src/nanoarrow/nanoarrow_hpp_test.cc
index 373c182f..4bbb8e9a 100644
--- a/src/nanoarrow/nanoarrow_hpp_test.cc
+++ b/src/nanoarrow/nanoarrow_hpp_test.cc
@@ -167,6 +167,70 @@ TEST(NanoarrowHppTest, NanoarrowHppUniqueBitmapTest) {
EXPECT_EQ(bitmap3->size_bits, 123);
}
+struct TestWrappedObj {
+ int64_t* num_frees;
+
+ TestWrappedObj(int64_t* addr) { num_frees = addr; }
+
+ TestWrappedObj(TestWrappedObj&& obj) {
+ num_frees = obj.num_frees;
+ obj.num_frees = nullptr;
+ }
+
+ ~TestWrappedObj() {
+ if (num_frees != nullptr) {
+ *num_frees = *num_frees + 1;
+ }
+ }
+};
+
+TEST(NanoarrowHppTest, NanoarrowHppBufferInitWrappedTest) {
+ nanoarrow::UniqueBuffer buffer;
+ int64_t num_frees = 0;
+
+ TestWrappedObj obj(&num_frees);
+ nanoarrow::BufferInitWrapped(buffer.get(), std::move(obj), nullptr, 0);
+ EXPECT_EQ(obj.num_frees, nullptr);
+ EXPECT_EQ(num_frees, 0);
+ buffer.reset();
+ EXPECT_EQ(num_frees, 1);
+
+ // Ensure the destructor won't get called again when ArrowBufferReset is
+ // called on the empty buffer.
+ buffer.reset();
+ EXPECT_EQ(num_frees, 1);
+}
+
+TEST(NanoarrowHppTest, NanoarrowHppBufferInitSequenceTest) {
+ nanoarrow::UniqueBuffer buffer;
+
+ // Check templating magic with std::string
+ nanoarrow::BufferInitSequence(buffer.get(), std::string("1234"));
+ EXPECT_EQ(buffer->size_bytes, 4);
+ EXPECT_EQ(buffer->capacity_bytes, 0);
+ EXPECT_EQ(memcmp(buffer->data, "1234", 4), 0);
+
+ // Check templating magic with std::vector
+ buffer.reset();
+ nanoarrow::BufferInitSequence(buffer.get(), std::vector<uint8_t>({1, 2, 3, 4}));
+ EXPECT_EQ(buffer->size_bytes, 4);
+ EXPECT_EQ(buffer->capacity_bytes, 0);
+ EXPECT_EQ(buffer->data[0], 1);
+ EXPECT_EQ(buffer->data[1], 2);
+ EXPECT_EQ(buffer->data[2], 3);
+ EXPECT_EQ(buffer->data[3], 4);
+
+ // Check templating magic with std::array
+ buffer.reset();
+ nanoarrow::BufferInitSequence(buffer.get(), std::array<uint8_t, 4>({1, 2, 3, 4}));
+ EXPECT_EQ(buffer->size_bytes, 4);
+ EXPECT_EQ(buffer->capacity_bytes, 0);
+ EXPECT_EQ(buffer->data[0], 1);
+ EXPECT_EQ(buffer->data[1], 2);
+ EXPECT_EQ(buffer->data[2], 3);
+ EXPECT_EQ(buffer->data[3], 4);
+}
+
TEST(NanoarrowHppTest, NanoarrowHppUniqueArrayViewTest) {
nanoarrow::UniqueArrayView array_view;
EXPECT_EQ(array_view->storage_type, NANOARROW_TYPE_UNINITIALIZED);
diff --git a/src/nanoarrow/utils.c b/src/nanoarrow/utils.c
index da4c4931..37ea6f62 100644
--- a/src/nanoarrow/utils.c
+++ b/src/nanoarrow/utils.c
@@ -201,7 +201,9 @@ static void ArrowBufferAllocatorMallocFree(struct ArrowBufferAllocator* allocato
uint8_t* ptr, int64_t size) {
NANOARROW_UNUSED(allocator);
NANOARROW_UNUSED(size);
- ArrowFree(ptr);
+ if (ptr != NULL) {
+ ArrowFree(ptr);
+ }
}
static struct ArrowBufferAllocator ArrowBufferAllocatorMalloc = {
@@ -211,13 +213,24 @@ struct ArrowBufferAllocator ArrowBufferAllocatorDefault(void) {
return ArrowBufferAllocatorMalloc;
}
-static uint8_t* ArrowBufferAllocatorNeverReallocate(
- struct ArrowBufferAllocator* allocator, uint8_t* ptr, int64_t old_size,
- int64_t new_size) {
- NANOARROW_UNUSED(allocator);
- NANOARROW_UNUSED(ptr);
- NANOARROW_UNUSED(old_size);
+static uint8_t* ArrowBufferDeallocatorReallocate(struct ArrowBufferAllocator* allocator,
+ uint8_t* ptr, int64_t old_size,
+ int64_t new_size) {
NANOARROW_UNUSED(new_size);
+
+ // Attempting to reallocate a buffer with a custom deallocator is
+ // a programming error. In debug mode, crash here.
+#if defined(NANOARROW_DEBUG)
+ NANOARROW_PRINT_AND_DIE(ENOMEM,
+ "It is an error to reallocate a buffer whose allocator is "
+ "ArrowBufferDeallocator()");
+#endif
+
+ // In release mode, ensure the the deallocator is called exactly
+ // once using the pointer it was given and return NULL, which
+ // will trigger the caller to return ENOMEM.
+ allocator->free(allocator, ptr, old_size);
+ *allocator = ArrowBufferAllocatorDefault();
return NULL;
}
@@ -226,7 +239,7 @@ struct ArrowBufferAllocator ArrowBufferDeallocator(
int64_t size),
void* private_data) {
struct ArrowBufferAllocator allocator;
- allocator.reallocate = &ArrowBufferAllocatorNeverReallocate;
+ allocator.reallocate = &ArrowBufferDeallocatorReallocate;
allocator.free = custom_free;
allocator.private_data = private_data;
return allocator;
diff --git a/src/nanoarrow/utils_test.cc b/src/nanoarrow/utils_test.cc
index caa4df57..8fca03b7 100644
--- a/src/nanoarrow/utils_test.cc
+++ b/src/nanoarrow/utils_test.cc
@@ -18,7 +18,6 @@
#include <cstring>
#include <string>
-#include <arrow/memory_pool.h>
#include <arrow/util/decimal.h>
#include <gtest/gtest.h>
@@ -97,28 +96,44 @@ TEST(ErrorTest, ErrorTestAssertNotOkDebug) {
TEST(ErrorTest, ErrorTestAssertNotOkRelease) { NANOARROW_ASSERT_OK(EINVAL); }
#endif
+class CustomMemoryPool {
+ public:
+ CustomMemoryPool() : bytes_allocated(0), num_reallocations(0), num_frees(0) {}
+ int64_t bytes_allocated;
+ int64_t num_reallocations;
+ int64_t num_frees;
+
+ static CustomMemoryPool* GetInstance() {
+ static CustomMemoryPool pool;
+ return &pool;
+ }
+};
+
static uint8_t* MemoryPoolReallocate(struct ArrowBufferAllocator* allocator, uint8_t* ptr,
int64_t old_size, int64_t new_size) {
- MemoryPool* pool = reinterpret_cast<MemoryPool*>(allocator->private_data);
- uint8_t* out = ptr;
- if (pool->Reallocate(old_size, new_size, &out).ok()) {
- return out;
- } else {
- return nullptr;
+ auto pool = reinterpret_cast<CustomMemoryPool*>(allocator->private_data);
+ pool->bytes_allocated -= old_size;
+ pool->num_reallocations += 1;
+
+ void* out = ArrowRealloc(ptr, new_size);
+ if (out != nullptr) {
+ pool->bytes_allocated += new_size;
}
+
+ return reinterpret_cast<uint8_t*>(out);
}
static void MemoryPoolFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr,
int64_t size) {
- MemoryPool* pool = reinterpret_cast<MemoryPool*>(allocator->private_data);
- pool->Free(ptr, size);
+ auto pool = reinterpret_cast<CustomMemoryPool*>(allocator->private_data);
+ pool->bytes_allocated -= size;
+ ArrowFree(ptr);
}
-static void MemoryPoolAllocatorInit(MemoryPool* pool,
- struct ArrowBufferAllocator* allocator) {
+static void MemoryPoolAllocatorInit(struct ArrowBufferAllocator* allocator) {
allocator->reallocate = &MemoryPoolReallocate;
allocator->free = &MemoryPoolFree;
- allocator->private_data = pool;
+ allocator->private_data = CustomMemoryPool::GetInstance();
}
TEST(AllocatorTest, AllocatorTestDefault) {
@@ -146,6 +161,7 @@ TEST(AllocatorTest, AllocatorTestDefault) {
// that keeps the buffer from being garbage collected (e.g., an SEXP in R)
struct CustomFreeData {
void* pointer_proxy;
+ int64_t num_frees;
};
static void CustomFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr,
@@ -153,39 +169,71 @@ static void CustomFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr,
auto data = reinterpret_cast<struct CustomFreeData*>(allocator->private_data);
ArrowFree(data->pointer_proxy);
data->pointer_proxy = nullptr;
+ data->num_frees++;
}
TEST(AllocatorTest, AllocatorTestDeallocator) {
struct CustomFreeData data;
- data.pointer_proxy = reinterpret_cast<uint8_t*>(ArrowMalloc(12));
+ data.pointer_proxy = nullptr;
+ data.num_frees = 0;
- struct ArrowBufferAllocator deallocator = ArrowBufferDeallocator(&CustomFree, &data);
+ struct ArrowBuffer buffer;
+ ArrowBufferInit(&buffer);
+
+ // Check that a zero-size, nullptr buffer is freed exactly once
+ ASSERT_EQ(ArrowBufferSetAllocator(&buffer, ArrowBufferDeallocator(&CustomFree, &data)),
+ NANOARROW_OK);
+ ASSERT_EQ(buffer.allocator.private_data, &data);
- EXPECT_EQ(deallocator.reallocate(&deallocator, nullptr, 0, 12), nullptr);
- EXPECT_EQ(deallocator.reallocate(&deallocator, nullptr, 0, 12), nullptr);
- deallocator.free(&deallocator, nullptr, 12);
- EXPECT_EQ(data.pointer_proxy, nullptr);
+ ArrowBufferReset(&buffer);
+ EXPECT_EQ(data.num_frees, 1);
+ EXPECT_NE(buffer.allocator.private_data, &data);
+
+ // Check reallocate behaviour
+ ASSERT_EQ(ArrowBufferSetAllocator(&buffer, ArrowBufferDeallocator(&CustomFree, &data)),
+ NANOARROW_OK);
+ data.num_frees = 0;
+
+ // In debug mode, an attempt to reallocate will crash with a message;
+ // in release mode it should result in ENOMEM.
+#if !defined(NANOARROW_DEBUG)
+ ASSERT_EQ(ArrowBufferAppendUInt8(&buffer, 0), ENOMEM);
+#endif
+ // Either way we should get exactly one free
+ ArrowBufferReset(&buffer);
+ EXPECT_EQ(data.num_frees, 1);
+
+ // Check wrapping an actual buffer
+ ASSERT_EQ(ArrowBufferSetAllocator(&buffer, ArrowBufferDeallocator(&CustomFree, &data)),
+ NANOARROW_OK);
+ buffer.data = reinterpret_cast<uint8_t*>(ArrowMalloc(12));
+ buffer.size_bytes = 12;
+ data.pointer_proxy = buffer.data;
+ data.num_frees = 0;
+
+ ArrowBufferReset(&buffer);
+ EXPECT_EQ(data.num_frees, 1);
}
TEST(AllocatorTest, AllocatorTestMemoryPool) {
struct ArrowBufferAllocator arrow_allocator;
- MemoryPoolAllocatorInit(system_memory_pool(), &arrow_allocator);
+ MemoryPoolAllocatorInit(&arrow_allocator);
- int64_t allocated0 = system_memory_pool()->bytes_allocated();
+ int64_t allocated0 = CustomMemoryPool::GetInstance()->bytes_allocated;
uint8_t* buffer = arrow_allocator.reallocate(&arrow_allocator, nullptr, 0, 10);
- EXPECT_EQ(system_memory_pool()->bytes_allocated() - allocated0, 10);
+ EXPECT_EQ(CustomMemoryPool::GetInstance()->bytes_allocated - allocated0, 10);
memset(buffer, 0, 10);
const char* test_str = "abcdefg";
memcpy(buffer, test_str, strlen(test_str) + 1);
buffer = arrow_allocator.reallocate(&arrow_allocator, buffer, 10, 100);
- EXPECT_EQ(system_memory_pool()->bytes_allocated() - allocated0, 100);
+ EXPECT_EQ(CustomMemoryPool::GetInstance()->bytes_allocated - allocated0, 100);
EXPECT_STREQ(reinterpret_cast<const char*>(buffer), test_str);
arrow_allocator.free(&arrow_allocator, buffer, 100);
- EXPECT_EQ(system_memory_pool()->bytes_allocated(), allocated0);
+ EXPECT_EQ(CustomMemoryPool::GetInstance()->bytes_allocated, allocated0);
buffer = arrow_allocator.reallocate(&arrow_allocator, nullptr, 0,
std::numeric_limits<int64_t>::max());