zhztheplayer commented on a change in pull request #7030:
URL: https://github.com/apache/arrow/pull/7030#discussion_r551694096
##########
File path: cpp/src/arrow/memory_pool.cc
##########
@@ -534,4 +535,139 @@ int64_t ProxyMemoryPool::max_memory() const { return
impl_->max_memory(); }
std::string ProxyMemoryPool::backend_name() const { return
impl_->backend_name(); }
+ReservationListener::~ReservationListener() {}
+
+ReservationListener::ReservationListener() {}
+
+class ReservationListenableMemoryPool::ReservationListenableMemoryPoolImpl {
+ public:
+ explicit ReservationListenableMemoryPoolImpl(
+ MemoryPool* pool, std::shared_ptr<ReservationListener> listener, int64_t
block_size)
+ : pool_(pool),
+ listener_(listener),
+ block_size_(block_size),
+ blocks_reserved_(0),
+ bytes_reserved_(0) {}
+
+ Status Allocate(int64_t size, uint8_t** out) {
+ RETURN_NOT_OK(UpdateReservation(size));
+ Status error = pool_->Allocate(size, out);
+ if (!error.ok()) {
+ RETURN_NOT_OK(UpdateReservation(-size));
+ return error;
+ }
+ return Status::OK();
+ }
+
+ Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) {
+ bool reserved = false;
+ int64_t diff = new_size - old_size;
+ if (new_size >= old_size) {
+ RETURN_NOT_OK(UpdateReservation(diff));
+ reserved = true;
Review comment:
Done, thanks for the suggestion.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
