From eeed82588bc9553dfccd463055d75ab6f6a06688 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 31 Mar 2026 03:55:44 +0800 Subject: [PATCH 1/3] fix --- cpp/core/memory/MemoryAllocator.cc | 2 +- cpp/core/tests/CMakeLists.txt | 1 + cpp/core/tests/MemoryAllocatorTest.cc | 83 +++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 cpp/core/tests/MemoryAllocatorTest.cc diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc index 0ca27085b22f..8e701c5145fb 100644 --- a/cpp/core/memory/MemoryAllocator.cc +++ b/cpp/core/memory/MemoryAllocator.cc @@ -137,7 +137,7 @@ bool StdMemoryAllocator::allocateZeroFilled(int64_t nmemb, int64_t size, void** if (*out == nullptr) { return false; } - bytes_ += size; + bytes_ += nmemb * size; return true; } diff --git a/cpp/core/tests/CMakeLists.txt b/cpp/core/tests/CMakeLists.txt index ac3c719db262..5bd34c77360d 100644 --- a/cpp/core/tests/CMakeLists.txt +++ b/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) diff --git a/cpp/core/tests/MemoryAllocatorTest.cc b/cpp/core/tests/MemoryAllocatorTest.cc new file mode 100644 index 000000000000..0e026f98ce9d --- /dev/null +++ b/cpp/core/tests/MemoryAllocatorTest.cc @@ -0,0 +1,83 @@ +/* + * 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 + +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. + void* buf = nullptr; + bool ok = allocator.allocateZeroFilled(10, 64, &buf); + ASSERT_TRUE(ok); + ASSERT_NE(buf, nullptr); + ASSERT_EQ(allocator.getBytes(), 640); + + // Free should bring bytes back to zero. + allocator.free(buf, 640); + ASSERT_EQ(allocator.getBytes(), 0); +} + +TEST(StdMemoryAllocator, allocateZeroFilledSingleElement) { + StdMemoryAllocator allocator; + + // nmemb=1 case: should track exactly size bytes. + void* buf = nullptr; + bool ok = allocator.allocateZeroFilled(1, 128, &buf); + ASSERT_TRUE(ok); + ASSERT_EQ(allocator.getBytes(), 128); + + allocator.free(buf, 128); + ASSERT_EQ(allocator.getBytes(), 0); +} + +TEST(StdMemoryAllocator, allocateZeroFilledZeroSize) { + StdMemoryAllocator allocator; + + // nmemb=0 or size=0: calloc returns non-null or null depending on impl, + // but bytes tracked should be 0 regardless. + void* buf = nullptr; + bool ok = allocator.allocateZeroFilled(0, 64, &buf); + if (ok) { + ASSERT_EQ(allocator.getBytes(), 0); + allocator.free(buf, 0); + } +} + +TEST(StdMemoryAllocator, allocateAndFreeAccounting) { + StdMemoryAllocator allocator; + + void* buf1 = nullptr; + void* buf2 = nullptr; + bool ok1 = allocator.allocate(100, &buf1); + ASSERT_TRUE(ok1); + ASSERT_EQ(allocator.getBytes(), 100); + + bool ok2 = allocator.allocateZeroFilled(5, 200, &buf2); + ASSERT_TRUE(ok2); + ASSERT_EQ(allocator.getBytes(), 1100); // 100 + 5*200 + + allocator.free(buf1, 100); + ASSERT_EQ(allocator.getBytes(), 1000); + + allocator.free(buf2, 1000); + ASSERT_EQ(allocator.getBytes(), 0); +} From 91df0374212c7efb45e3d33e3a3cfa8ffe10fe68 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 31 Mar 2026 18:37:47 +0800 Subject: [PATCH 2/3] [VL] Trim MemoryAllocatorTest to single test case for allocateZeroFilled fix Keep only the test that directly validates the nmemb * size accounting bug; remove three unrelated general-purpose allocator tests. --- cpp/core/tests/MemoryAllocatorTest.cc | 47 --------------------------- 1 file changed, 47 deletions(-) diff --git a/cpp/core/tests/MemoryAllocatorTest.cc b/cpp/core/tests/MemoryAllocatorTest.cc index 0e026f98ce9d..1abec4be5c66 100644 --- a/cpp/core/tests/MemoryAllocatorTest.cc +++ b/cpp/core/tests/MemoryAllocatorTest.cc @@ -31,53 +31,6 @@ TEST(StdMemoryAllocator, allocateZeroFilledAccounting) { ASSERT_NE(buf, nullptr); ASSERT_EQ(allocator.getBytes(), 640); - // Free should bring bytes back to zero. allocator.free(buf, 640); ASSERT_EQ(allocator.getBytes(), 0); } - -TEST(StdMemoryAllocator, allocateZeroFilledSingleElement) { - StdMemoryAllocator allocator; - - // nmemb=1 case: should track exactly size bytes. - void* buf = nullptr; - bool ok = allocator.allocateZeroFilled(1, 128, &buf); - ASSERT_TRUE(ok); - ASSERT_EQ(allocator.getBytes(), 128); - - allocator.free(buf, 128); - ASSERT_EQ(allocator.getBytes(), 0); -} - -TEST(StdMemoryAllocator, allocateZeroFilledZeroSize) { - StdMemoryAllocator allocator; - - // nmemb=0 or size=0: calloc returns non-null or null depending on impl, - // but bytes tracked should be 0 regardless. - void* buf = nullptr; - bool ok = allocator.allocateZeroFilled(0, 64, &buf); - if (ok) { - ASSERT_EQ(allocator.getBytes(), 0); - allocator.free(buf, 0); - } -} - -TEST(StdMemoryAllocator, allocateAndFreeAccounting) { - StdMemoryAllocator allocator; - - void* buf1 = nullptr; - void* buf2 = nullptr; - bool ok1 = allocator.allocate(100, &buf1); - ASSERT_TRUE(ok1); - ASSERT_EQ(allocator.getBytes(), 100); - - bool ok2 = allocator.allocateZeroFilled(5, 200, &buf2); - ASSERT_TRUE(ok2); - ASSERT_EQ(allocator.getBytes(), 1100); // 100 + 5*200 - - allocator.free(buf1, 100); - ASSERT_EQ(allocator.getBytes(), 1000); - - allocator.free(buf2, 1000); - ASSERT_EQ(allocator.getBytes(), 0); -} From c7159a2bd92c676c73d447981ec39e90e7370222 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 31 Mar 2026 20:17:12 +0800 Subject: [PATCH 3/3] address comments --- cpp/core/memory/MemoryAllocator.cc | 5 +++++ cpp/core/tests/MemoryAllocatorTest.cc | 5 +++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc index 8e701c5145fb..c8aaa1385e8d 100644 --- a/cpp/core/memory/MemoryAllocator.cc +++ b/cpp/core/memory/MemoryAllocator.cc @@ -18,6 +18,8 @@ #include "MemoryAllocator.h" #include "utils/Macros.h" +#include + namespace gluten { bool ListenableMemoryAllocator::allocate(int64_t size, void** out) { @@ -133,6 +135,9 @@ 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::max() / size, + "nmemb * size overflows int64_t"); *out = std::calloc(nmemb, size); if (*out == nullptr) { return false; diff --git a/cpp/core/tests/MemoryAllocatorTest.cc b/cpp/core/tests/MemoryAllocatorTest.cc index 1abec4be5c66..92b337d8a6da 100644 --- a/cpp/core/tests/MemoryAllocatorTest.cc +++ b/cpp/core/tests/MemoryAllocatorTest.cc @@ -25,12 +25,13 @@ TEST(StdMemoryAllocator, allocateZeroFilledAccounting) { 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(), 640); + ASSERT_EQ(allocator.getBytes(), expectedBytes); - allocator.free(buf, 640); + ASSERT_TRUE(allocator.free(buf, expectedBytes)); ASSERT_EQ(allocator.getBytes(), 0); }