diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc index 0ca27085b22f..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,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::max() / size, + "nmemb * size overflows int64_t"); *out = std::calloc(nmemb, size); 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..92b337d8a6da --- /dev/null +++ b/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 + +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); +}