diff --git a/cpp/src/parquet/util/mem-allocator.cc b/cpp/src/parquet/util/mem-allocator.cc index 4d42bc4925024..2b6592d052493 100644 --- a/cpp/src/parquet/util/mem-allocator.cc +++ b/cpp/src/parquet/util/mem-allocator.cc @@ -30,17 +30,23 @@ uint8_t* TrackingAllocator::Malloc(int64_t size) { uint8_t* p = static_cast(std::malloc(size)); if (!p) { throw ParquetException("OOM: memory allocation failed"); } - total_memory_ += size; - if (total_memory_ > max_memory_) { max_memory_ = total_memory_; } + { + std::lock_guard lock(stats_mutex_); + total_memory_ += size; + if (total_memory_ > max_memory_) { max_memory_ = total_memory_; } + } return p; } void TrackingAllocator::Free(uint8_t* p, int64_t size) { if (nullptr != p && size > 0) { - if (total_memory_ < size) { - throw ParquetException("Attempting to free too much memory"); + { + std::lock_guard lock(stats_mutex_); + if (total_memory_ < size) { + throw ParquetException("Attempting to free too much memory"); + } + total_memory_ -= size; } - total_memory_ -= size; std::free(p); } } diff --git a/cpp/src/parquet/util/mem-allocator.h b/cpp/src/parquet/util/mem-allocator.h index b20d9f9ebc7b5..a0f3693307e2c 100644 --- a/cpp/src/parquet/util/mem-allocator.h +++ b/cpp/src/parquet/util/mem-allocator.h @@ -19,6 +19,7 @@ #define PARQUET_UTIL_MEMORY_POOL_H #include +#include #include "parquet/util/visibility.h" @@ -48,6 +49,7 @@ class PARQUET_EXPORT TrackingAllocator : public MemoryAllocator { int64_t MaxMemory() { return max_memory_; } private: + std::mutex stats_mutex_; int64_t total_memory_; int64_t max_memory_; };