1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17
18 #include <gtest/gtest.h>
19
20 #include "arrow/memory_pool.h"
21 #include "arrow/testing/gtest_util.h"
22 #include "jni/dataset/jni_util.h"
23
24 namespace arrow {
25 namespace dataset {
26 namespace jni {
27
28 class MyListener : public ReservationListener {
29 public:
OnReservation(int64_t size)30 Status OnReservation(int64_t size) override {
31 bytes_reserved_ += size;
32 reservation_count_++;
33 return arrow::Status::OK();
34 }
35
OnRelease(int64_t size)36 Status OnRelease(int64_t size) override {
37 bytes_reserved_ -= size;
38 release_count_++;
39 return arrow::Status::OK();
40 }
41
bytes_reserved()42 int64_t bytes_reserved() { return bytes_reserved_; }
43
reservation_count() const44 int32_t reservation_count() const { return reservation_count_; }
45
release_count() const46 int32_t release_count() const { return release_count_; }
47
48 private:
49 int64_t bytes_reserved_;
50 int32_t reservation_count_;
51 int32_t release_count_;
52 };
53
TEST(ReservationListenableMemoryPool,Basic)54 TEST(ReservationListenableMemoryPool, Basic) {
55 auto pool = MemoryPool::CreateDefault();
56 auto listener = std::make_shared<MyListener>();
57 ReservationListenableMemoryPool rlp(pool.get(), listener);
58
59 uint8_t* data;
60 ASSERT_OK(rlp.Allocate(100, &data));
61
62 uint8_t* data2;
63 ASSERT_OK(rlp.Allocate(100, &data2));
64
65 rlp.Free(data, 100);
66 rlp.Free(data2, 100);
67
68 ASSERT_EQ(200, rlp.max_memory());
69 ASSERT_EQ(200, pool->max_memory());
70 }
71
TEST(ReservationListenableMemoryPool,Listener)72 TEST(ReservationListenableMemoryPool, Listener) {
73 auto pool = MemoryPool::CreateDefault();
74 auto listener = std::make_shared<MyListener>();
75 ReservationListenableMemoryPool rlp(pool.get(), listener);
76
77 uint8_t* data;
78 ASSERT_OK(rlp.Allocate(100, &data));
79
80 uint8_t* data2;
81 ASSERT_OK(rlp.Allocate(100, &data2));
82
83 ASSERT_EQ(200, rlp.bytes_allocated());
84 ASSERT_EQ(512 * 1024, listener->bytes_reserved());
85
86 rlp.Free(data, 100);
87 rlp.Free(data2, 100);
88
89 ASSERT_EQ(0, rlp.bytes_allocated());
90 ASSERT_EQ(0, listener->bytes_reserved());
91 ASSERT_EQ(1, listener->reservation_count());
92 ASSERT_EQ(1, listener->release_count());
93 }
94
TEST(ReservationListenableMemoryPool,BlockSize)95 TEST(ReservationListenableMemoryPool, BlockSize) {
96 auto pool = MemoryPool::CreateDefault();
97 auto listener = std::make_shared<MyListener>();
98 ReservationListenableMemoryPool rlp(pool.get(), listener, 100);
99
100 uint8_t* data;
101 ASSERT_OK(rlp.Allocate(100, &data));
102
103 ASSERT_EQ(100, rlp.bytes_allocated());
104 ASSERT_EQ(100, listener->bytes_reserved());
105
106 rlp.Free(data, 100);
107
108 ASSERT_EQ(0, rlp.bytes_allocated());
109 ASSERT_EQ(0, listener->bytes_reserved());
110 }
111
TEST(ReservationListenableMemoryPool,BlockSize2)112 TEST(ReservationListenableMemoryPool, BlockSize2) {
113 auto pool = MemoryPool::CreateDefault();
114 auto listener = std::make_shared<MyListener>();
115 ReservationListenableMemoryPool rlp(pool.get(), listener, 99);
116
117 uint8_t* data;
118 ASSERT_OK(rlp.Allocate(100, &data));
119
120 ASSERT_EQ(100, rlp.bytes_allocated());
121 ASSERT_EQ(198, listener->bytes_reserved());
122
123 rlp.Free(data, 100);
124
125 ASSERT_EQ(0, rlp.bytes_allocated());
126 ASSERT_EQ(0, listener->bytes_reserved());
127
128 ASSERT_EQ(1, listener->reservation_count());
129 ASSERT_EQ(1, listener->release_count());
130 }
131
132 } // namespace jni
133 } // namespace dataset
134 } // namespace arrow
135