implementing feedback

Author: rok

cpp/src/arrow/memory_pool-test.cc

@@ -95,21 +95,21 @@ TEST(LoggingMemoryPool, Logging) {
 TEST(ProxyMemoryPool, Logging) {
   MemoryPool* pool = default_memory_pool();
 
-  ProxyMemoryPool pp = ProxyMemoryPool(pool);
+  ProxyMemoryPool* pp = new ProxyMemoryPool(pool);
 
   uint8_t* data;
   ASSERT_OK(pool->Allocate(100, &data));
 
   uint8_t* data2;
-  ASSERT_OK(pp.Allocate(300, &data2));
+  ASSERT_OK(pp->Allocate(300, &data2));
 
   ASSERT_EQ(400, pool->bytes_allocated());
-  ASSERT_EQ(300, pp.bytes_allocated());
+  ASSERT_EQ(300, pp->bytes_allocated());
 
   pool->Free(data, 100);
-  pp.Free(data2, 300);
+  pp->Free(data2, 300);
 
   ASSERT_EQ(0, pool->bytes_allocated());
-  ASSERT_EQ(0, pp.bytes_allocated());
+  ASSERT_EQ(0, pp->bytes_allocated());
 }
 }  // namespace arrow

cpp/src/arrow/memory_pool.cc

@@ -205,29 +205,36 @@ int64_t LoggingMemoryPool::max_memory() const {
 ProxyMemoryPool::ProxyMemoryPool(MemoryPool* pool) : pool_(pool) {}
 
 Status ProxyMemoryPool::Allocate(int64_t size, uint8_t** out) {
-  Status s = pool_->Allocate(size, out);
+  RETURN_NOT_OK(pool_->Allocate(size, out));
   bytes_allocated_ += size;
-  return s;
+  {
+    std::lock_guard<std::mutex> guard(lock_);
+    if (bytes_allocated_ > max_memory_) {
+      max_memory_ = bytes_allocated_.load();
+    }
+  }
+  return Status::OK();
 }
 
 Status ProxyMemoryPool::Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) {
-  Status s = pool_->Reallocate(old_size, new_size, ptr);
+  RETURN_NOT_OK(pool_->Reallocate(old_size, new_size, ptr));
   bytes_allocated_ += new_size - old_size;
-  return s;
+  {
+    std::lock_guard<std::mutex> guard(lock_);
+    if (bytes_allocated_ > max_memory_) {
+      max_memory_ = bytes_allocated_.load();
+    }
+  }
+  return Status::OK();
 }
 
 void ProxyMemoryPool::Free(uint8_t* buffer, int64_t size) {
   pool_->Free(buffer, size);
   bytes_allocated_ -= size;
 }
 
-int64_t ProxyMemoryPool::bytes_allocated() const {
-  int64_t nb_bytes = bytes_allocated_;
-  return nb_bytes;
-}
+int64_t ProxyMemoryPool::bytes_allocated() const { return bytes_allocated_.load(); }
+
+int64_t ProxyMemoryPool::max_memory() const { return max_memory_.load(); }
 
-int64_t ProxyMemoryPool::max_memory() const {
-  int64_t mem = pool_->max_memory();
-  return mem;
-}
 }  // namespace arrow

cpp/src/arrow/memory_pool.h

@@ -19,6 +19,7 @@
 #define ARROW_MEMORY_POOL_H
 
 #include <atomic>
+#include <mutex>
 #include <cstdint>
 
 #include "arrow/util/visibility.h"
@@ -86,10 +87,13 @@ class ARROW_EXPORT LoggingMemoryPool : public MemoryPool {
   MemoryPool* pool_;
 };
 
+/// Derived class for memory allocation.
+///
+/// Tracks the number of bytes and maximum mmeory allocated through its direct
+/// calls. Actual allocation is delegated to MemoryPool class.
 class ARROW_EXPORT ProxyMemoryPool : public MemoryPool {
  public:
   explicit ProxyMemoryPool(MemoryPool* pool);
-  ~ProxyMemoryPool() override = default;
 
   Status Allocate(int64_t size, uint8_t** out) override;
   Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override;
@@ -101,8 +105,10 @@ class ARROW_EXPORT ProxyMemoryPool : public MemoryPool {
   int64_t max_memory() const override;
 
  private:
+  mutable std::mutex lock_;
   MemoryPool* pool_;
-  int64_t bytes_allocated_ = 0;
+  std::atomic<int64_t> bytes_allocated_{0};
+  std::atomic<int64_t> max_memory_{0};
 };
 
 ARROW_EXPORT MemoryPool* default_memory_pool();

python/pyarrow/includes/libarrow.pxd

@@ -194,7 +194,6 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     cdef cppclass CProxyMemoryPool" arrow::ProxyMemoryPool"(CMemoryPool):
         CProxyMemoryPool(CMemoryPool*)
-        int64_t bytes_allocated()
 
     cdef cppclass CBuffer" arrow::Buffer":
         CBuffer(const uint8_t* data, int64_t size)

python/pyarrow/memory.pxi

@@ -67,8 +67,6 @@ def set_memory_pool(MemoryPool pool):
 cdef MemoryPool _default_memory_pool = default_memory_pool()
 cdef LoggingMemoryPool _logging_memory_pool = (
     LoggingMemoryPool(_default_memory_pool))
-cdef ProxyMemoryPool _proxy_memory_pool = (
-    ProxyMemoryPool(_default_memory_pool))
 
 
 def log_memory_allocations(enable=True):