include/arrow/memory_pool_test.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <algorithm>
0021 #include <cstddef>
0022 #include <cstdint>
0023 #include <limits>
0024
0025 #include <gtest/gtest.h>
0026
0027 #include "arrow/memory_pool.h"
0028 #include "arrow/status.h"
0029 #include "arrow/testing/gtest_util.h"
0030
0031 namespace arrow {
0032
0033 class TestMemoryPoolBase : public ::testing::Test {
0034  public:
0035   virtual ::arrow::MemoryPool* memory_pool() = 0;
0036
0037   void TestMemoryTracking() {
0038     auto pool = memory_pool();
0039
0040     uint8_t* data;
0041     const auto old_bytes_allocated = pool->bytes_allocated();
0042     ASSERT_OK(pool->Allocate(100, &data));
0043     EXPECT_EQ(static_cast<uint64_t>(0), reinterpret_cast<uint64_t>(data) % 64);
0044     ASSERT_EQ(old_bytes_allocated + 100, pool->bytes_allocated());
0045
0046     uint8_t* data2;
0047     ASSERT_OK(pool->Allocate(27, &data2));
0048     EXPECT_EQ(static_cast<uint64_t>(0), reinterpret_cast<uint64_t>(data2) % 64);
0049     ASSERT_EQ(old_bytes_allocated + 127, pool->bytes_allocated());
0050
0051     pool->Free(data, 100);
0052     ASSERT_EQ(old_bytes_allocated + 27, pool->bytes_allocated());
0053     pool->Free(data2, 27);
0054     ASSERT_EQ(old_bytes_allocated, pool->bytes_allocated());
0055   }
0056
0057   void TestOOM() {
0058     auto pool = memory_pool();
0059
0060     uint8_t* data;
0061     int64_t max_alloc = std::min<uint64_t>(std::numeric_limits<int64_t>::max(),
0062                                            std::numeric_limits<size_t>::max());
0063     // subtract 63 to prevent overflow after the size is aligned
0064     for (int64_t to_alloc : {max_alloc, max_alloc - 63, max_alloc - 127}) {
0065       ASSERT_RAISES(OutOfMemory, pool->Allocate(to_alloc, &data));
0066     }
0067   }
0068
0069   void TestReallocate() {
0070     auto pool = memory_pool();
0071
0072     uint8_t* data;
0073     ASSERT_OK(pool->Allocate(10, &data));
0074     ASSERT_EQ(10, pool->bytes_allocated());
0075     data[0] = 35;
0076     data[9] = 12;
0077
0078     // Expand
0079     ASSERT_OK(pool->Reallocate(10, 20, &data));
0080     ASSERT_EQ(data[9], 12);
0081     ASSERT_EQ(20, pool->bytes_allocated());
0082
0083     // Shrink
0084     ASSERT_OK(pool->Reallocate(20, 5, &data));
0085     ASSERT_EQ(data[0], 35);
0086     ASSERT_EQ(5, pool->bytes_allocated());
0087
0088     // Free
0089     pool->Free(data, 5);
0090     ASSERT_EQ(0, pool->bytes_allocated());
0091   }
0092
0093   void TestAlignment() {
0094     auto pool = memory_pool();
0095     {
0096       uint8_t* data64;
0097       ASSERT_OK(pool->Allocate(10, &data64));
0098       ASSERT_EQ(reinterpret_cast<uintptr_t>(data64) % kDefaultBufferAlignment, 0);
0099       pool->Free(data64, 10);
0100     }
0101
0102     {
0103       uint8_t* data512;
0104       ASSERT_OK(pool->Allocate(10, 512, &data512));
0105       ASSERT_EQ(reinterpret_cast<uintptr_t>(data512) % 512, 0);
0106       pool->Free(data512, 10, 512);
0107     }
0108   }
0109 };
0110
0111 }  // namespace arrow