[VL] Fix StdMemoryAllocator::allocateZeroFilled byte accounting

cpp/core/memory/MemoryAllocator.cc

@@ -18,6 +18,8 @@
 #include "MemoryAllocator.h"
 #include "utils/Macros.h"
 
+#include <limits>
+
 namespace gluten {
 
 bool ListenableMemoryAllocator::allocate(int64_t size, void** out) {
@@ -133,11 +135,14 @@ bool StdMemoryAllocator::allocate(int64_t size, void** out) {
 bool StdMemoryAllocator::allocateZeroFilled(int64_t nmemb, int64_t size, void** out) {
   GLUTEN_CHECK(nmemb >= 0, "nmemb is less than 0");
   GLUTEN_CHECK(size >= 0, "size is less than 0");
+  GLUTEN_CHECK(
+      size == 0 || nmemb <= std::numeric_limits<int64_t>::max() / size,
+      "nmemb * size overflows int64_t");
   *out = std::calloc(nmemb, size);
   if (*out == nullptr) {
     return false;
   }
-  bytes_ += size;
+  bytes_ += nmemb * size;
   return true;
 }
 

cpp/core/tests/CMakeLists.txt

@@ -15,3 +15,4 @@
 
 add_test_case(round_robin_partitioner_test SOURCES RoundRobinPartitionerTest.cc)
 add_test_case(object_store_test SOURCES ObjectStoreTest.cc)
+add_test_case(memory_allocator_test SOURCES MemoryAllocatorTest.cc)

cpp/core/tests/MemoryAllocatorTest.cc

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "memory/MemoryAllocator.h"
+#include <gtest/gtest.h>
+
+using namespace gluten;
+
+TEST(StdMemoryAllocator, allocateZeroFilledAccounting) {
+  StdMemoryAllocator allocator;
+  ASSERT_EQ(allocator.getBytes(), 0);
+
+  // allocateZeroFilled with nmemb=10, size=64 should track 10*64=640 bytes.
+  const int64_t expectedBytes = 10 * 64;
+  void* buf = nullptr;
+  bool ok = allocator.allocateZeroFilled(10, 64, &buf);
+  ASSERT_TRUE(ok);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(allocator.getBytes(), expectedBytes);
+
+  ASSERT_TRUE(allocator.free(buf, expectedBytes));
+  ASSERT_EQ(allocator.getBytes(), 0);
+}