Memory leak in SwissTable::grow_double()

[zifengyu]

Describe the bug, including details regarding any error messages, version, and platform.

There are two allocation in grow_double. If the first suceeded but the second failed, the function will return without Free the first buffer

RETURN_NOT_OK(pool_->Allocate(block_size_total_after, &blocks_new));
RETURN_NOT_OK(pool_->Allocate(hashes_size_total_after, &hashes_new_8B));

https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/key_map.cc#L670

Component(s)

C++

[mapleFU]

I think pool_->Allocate returns a buffer with MemoryManager:

  static std::unique_ptr<PoolBuffer> MakeUnique(MemoryPool* pool, int64_t alignment) {
    std::shared_ptr<MemoryManager> mm;
    if (pool == nullptr) {
      pool = default_memory_pool();
      mm = default_cpu_memory_manager();
    } else {
      mm = CPUDevice::memory_manager(pool);
    }
    return std::make_unique<PoolBuffer>(std::move(mm), pool, alignment);
  }

When PoolBuffer destruct, it will release pool memory:

  ~PoolBuffer() override {
    // Avoid calling pool_->Free if the global pools are destroyed
    // (XXX this will not work with user-defined pools)

    // This can happen if a Future is destructing on one thread while or
    // after memory pools are destructed on the main thread (as there is
    // no guarantee of destructor order between thread/memory pools)
    uint8_t* ptr = mutable_data();
    if (ptr && !global_state.is_finalizing()) {
      pool_->Free(ptr, capacity_, alignment_);
    }
  }

If I'm wrong, please point out that

[zifengyu]

Hi @mapleFU,

SwissTable managers the pool memory directly instead of using PoolBuffer. If the grow_doulbe quit in the middle, the blocks_new is not saved to blocks_, the explict release memory/Free is not called.
https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/key_map.cc#L770

[mapleFU]

Yes, you're right, seems this may causing memory leak. @westonpace Mind take a look?
@zifengyu Or would you mind give a quick fixing?

[westonpace]

Thanks for the catch. I think smart pointers should be preferred in general and I've added a PR that switches these two buffers to using Buffer.