Add guard back

cpp/src/arrow/compute/util_internal.cc

@@ -38,7 +38,7 @@ TempVectorStack::~TempVectorStack() {
 Status TempVectorStack::Init(MemoryPool* pool, int64_t size) {
   num_vectors_ = 0;
   top_ = 0;
-  buffer_size_ = PaddedAllocationSize(size);
+  buffer_size_ = EstimatedAllocationSize(size);
   ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateResizableBuffer(size, pool));
 #ifdef ADDRESS_SANITIZER
   ASAN_POISON_MEMORY_REGION(buffer->mutable_data(), size);
@@ -59,26 +59,37 @@ int64_t TempVectorStack::PaddedAllocationSize(int64_t num_bytes) {
 }
 
 void TempVectorStack::alloc(uint32_t num_bytes, uint8_t** data, int* id) {
-  int64_t alloc_size = PaddedAllocationSize(num_bytes);
-  int64_t new_top = top_ + alloc_size;
+  int64_t estimated_alloc_size = EstimatedAllocationSize(num_bytes);
+  int64_t new_top = top_ + estimated_alloc_size;
   // Stack overflow check (see GH-39582).
   // XXX cannot return a regular Status because most consumers do not either.
   ARROW_CHECK_LE(new_top, buffer_size_)
-      << "TempVectorStack::alloc overflow: allocating " << alloc_size << " on top of "
-      << top_ << " in stack of size " << buffer_size_;
-  *data = buffer_->mutable_data() + top_;
+      << "TempVectorStack::alloc overflow: allocating " << estimated_alloc_size
+      << " on top of " << top_ << " in stack of size " << buffer_size_;
 #ifdef ADDRESS_SANITIZER
-  ASAN_UNPOISON_MEMORY_REGION(*data, alloc_size);
+  ASAN_UNPOISON_MEMORY_REGION(buffer_->mutable_data() + top_, estimated_alloc_size);
+#endif
+  *data = buffer_->mutable_data() + top_ + /*one guard*/ sizeof(uint64_t);
+#ifndef NDEBUG
+  // We set 8 bytes before the beginning of the allocated range and
+  // 8 bytes after the end to check for stack overflow (which would
+  // result in those known bytes being corrupted).
+  reinterpret_cast<uint64_t*>(buffer_->mutable_data() + top_)[0] = kGuard1;
+  reinterpret_cast<uint64_t*>(buffer_->mutable_data() + new_top)[-1] = kGuard2;
 #endif
   *id = num_vectors_++;
   top_ = new_top;
 }
 
 void TempVectorStack::release(int id, uint32_t num_bytes) {
   ARROW_DCHECK(num_vectors_ == id + 1);
-  int64_t size = PaddedAllocationSize(num_bytes);
+  int64_t size = EstimatedAllocationSize(num_bytes);
+  ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[-1] ==
+               kGuard2);
   ARROW_DCHECK(top_ >= size);
   top_ -= size;
+  ARROW_DCHECK(reinterpret_cast<const uint64_t*>(buffer_->mutable_data() + top_)[0] ==
+               kGuard1);
 #ifdef ADDRESS_SANITIZER
   ASAN_POISON_MEMORY_REGION(buffer_->mutable_data() + top_, size);
 #endif

cpp/src/arrow/compute/util_internal.h

@@ -51,10 +51,16 @@ class ARROW_EXPORT TempVectorStack {
   int64_t AllocatedSize() const { return top_; }
 
  private:
+  static int64_t EstimatedAllocationSize(int64_t size) {
+    return PaddedAllocationSize(size) + /*two guards*/ 2 * sizeof(uint64_t);
+  }
+
   static int64_t PaddedAllocationSize(int64_t num_bytes);
 
   void alloc(uint32_t num_bytes, uint8_t** data, int* id);
   void release(int id, uint32_t num_bytes);
+  static constexpr uint64_t kGuard1 = 0x3141592653589793ULL;
+  static constexpr uint64_t kGuard2 = 0x0577215664901532ULL;
   static constexpr int64_t kPadding = 64;
   int num_vectors_;
   int64_t top_;