jvanstraten commented on a change in pull request #12116:
URL: https://github.com/apache/arrow/pull/12116#discussion_r807164882
##########
File path: cpp/src/arrow/memory_pool.cc
##########
@@ -603,14 +643,116 @@ class BaseMemoryPoolImpl : public MemoryPool {
stats_.UpdateAllocatedBytes(-size);
}
- void ReleaseUnused() override { Allocator::ReleaseUnused(); }
+ protected:
+ virtual Status AllocateImmutableZeros(int64_t size, uint8_t** out) {
+#ifdef USE_MMAP_FOR_IMMUTABLE_ZEROS
+ if (size > 0) {
+ *out = static_cast<uint8_t*>(mmap(
+ nullptr, size, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS |
MAP_NORESERVE, -1, 0));
+ if (*out == MAP_FAILED) {
+ auto err = errno;
+ return Status::OutOfMemory("Failed to allocate zero buffer of size ",
size, ": ",
+ strerror(err));
+ }
+ return Status::OK();
+ }
+#endif
+ // TODO: jemalloc and mimalloc support zero-initialized allocations as
+ // well, which might be faster than allocate + memset.
+ RETURN_NOT_OK(Allocate(size, out));
+ std::memset(*out, 0, size);
+ return Status::OK();
+ }
+
+ void FreeImmutableZeros(uint8_t* buffer, int64_t size) override {
+#ifdef USE_MMAP_FOR_IMMUTABLE_ZEROS
+ if (size > 0) {
+ munmap(buffer, size);
+ return;
+ }
+#endif
+ Free(buffer, size);
+ }
+
+ public:
+ Result<std::shared_ptr<Buffer>> GetImmutableZeros(int64_t size) override {
+ // Thread-safely get the current largest buffer of zeros.
+ auto current_buffer = atomic_load(&immutable_zeros_cache_);
+
+ // If this buffer satisfies the requirements, return it.
+ if (current_buffer && current_buffer->size() >= size) {
+ return std::move(current_buffer);
+ }
+
+ // Acquire the lock for allocating a new buffer.
+ std::lock_guard<std::mutex> gg(immutable_zeros_mutex_);
+
+ // Between our previous atomic load and acquisition of the lock, another
+ // thread may have allocated a buffer. So we need to check again.
+ current_buffer = atomic_load(&immutable_zeros_cache_);
+ if (current_buffer && current_buffer->size() >= size) {
+ return std::move(current_buffer);
+ }
+
+ // Let's now figure out a good size to allocate. This is done
+ // heuristically, with the following rules:
+ // - allocate at least the requested size (obviously);
+ // - allocate at least 2x the previous size;
+ // - allocate at least kMinAllocSize bytes (to avoid lots of small
+ // allocations).
+ static const int64_t kMinAllocSize = 4096;
+ int64_t alloc_size =
+ std::max(size, current_buffer ? (current_buffer->size() * 2) :
kMinAllocSize);
+
+ // Attempt to allocate the block.
+ uint8_t* data = nullptr;
+ auto result = AllocateImmutableZeros(alloc_size, &data);
+
+ // If we fail to do so, fall back to trying to allocate the requested size
+ // exactly as a last-ditch effort.
+ if (!result.ok() || data == nullptr) {
+ alloc_size = size;
+ RETURN_NOT_OK(AllocateImmutableZeros(alloc_size, &data));
+ }
+ DCHECK_NE(data, nullptr);
+
+ // Move ownership of the data block into an ImmutableZeros object. It will
+ // free the block when destroyed, i.e. when all shared_ptr references to it
+ // are reset or go out of scope.
+ current_buffer = std::make_shared<ImmutableZeros>(data, alloc_size, this);
+
+ // Store a reference to the new block in the cache, so subsequent calls to
+ // this function (from this thread or from other threads) can use it, too.
+ atomic_store(&immutable_zeros_cache_, current_buffer);
+
+ return std::move(current_buffer);
+ }
+
+ void ReleaseUnused() override {
+ // Get rid of the ImmutableZeros cache if we're the only one using it. If
+ // there are other pieces of code using it, getting rid of the cache won't
+ // deallocate it anyway, so it's better to hold onto it.
+ {
+ auto cache = atomic_load(&immutable_zeros_cache_);
+
+ // Because we now have a copy in our thread, the use count will be 2 if
+ // nothing else is using it.
+ if (cache.use_count() <= 2) {
+ atomic_store(&immutable_zeros_cache_,
std::shared_ptr<ImmutableZeros>());
Review comment:
It's not entirely thread-safe in the sense that if other threads are
doing stuff with the cache while `ReleaseUnused()` is called, `ReleaseUnused()`
may or may not flush the cache due to a data race in `use_count`. But it
doesn't break anything in either case, and `ReleaseUnused()` is already
documented to be best-effort.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
