1 // Protocol Buffers - Google's data interchange format
2 // Copyright 2023 Google LLC. All rights reserved.
3 //
4 // Use of this source code is governed by a BSD-style
5 // license that can be found in the LICENSE file or at
6 // https://developers.google.com/open-source/licenses/bsd
7
8 #include "upb/mem/arena.h"
9
10 #include <stddef.h>
11
12 #include <array>
13 #include <atomic>
14 #include <thread>
15 #include <vector>
16
17 #include <gtest/gtest.h>
18 #include "absl/random/distributions.h"
19 #include "absl/random/random.h"
20 #include "absl/synchronization/notification.h"
21 #include "absl/time/clock.h"
22 #include "absl/time/time.h"
23 #include "upb/mem/alloc.h"
24
25 // Must be last.
26 #include "upb/port/def.inc"
27
28 namespace {
29
TEST(ArenaTest,ArenaFuse)30 TEST(ArenaTest, ArenaFuse) {
31 upb_Arena* arena1 = upb_Arena_New();
32 upb_Arena* arena2 = upb_Arena_New();
33
34 EXPECT_TRUE(upb_Arena_Fuse(arena1, arena2));
35
36 upb_Arena_Free(arena1);
37 upb_Arena_Free(arena2);
38 }
39
40 // Do-nothing allocator for testing.
TestAllocFunc(upb_alloc * alloc,void * ptr,size_t oldsize,size_t size)41 extern "C" void* TestAllocFunc(upb_alloc* alloc, void* ptr, size_t oldsize,
42 size_t size) {
43 return upb_alloc_global.func(alloc, ptr, oldsize, size);
44 }
45
TEST(ArenaTest,FuseWithInitialBlock)46 TEST(ArenaTest, FuseWithInitialBlock) {
47 char buf1[1024];
48 char buf2[1024];
49 upb_Arena* arenas[] = {upb_Arena_Init(buf1, 1024, &upb_alloc_global),
50 upb_Arena_Init(buf2, 1024, &upb_alloc_global),
51 upb_Arena_Init(nullptr, 0, &upb_alloc_global)};
52 int size = sizeof(arenas) / sizeof(arenas[0]);
53 for (int i = 0; i < size; ++i) {
54 for (int j = 0; j < size; ++j) {
55 if (i == j) {
56 // Fuse to self is always allowed.
57 EXPECT_TRUE(upb_Arena_Fuse(arenas[i], arenas[j]));
58 } else {
59 EXPECT_FALSE(upb_Arena_Fuse(arenas[i], arenas[j]));
60 }
61 }
62 }
63
64 for (int i = 0; i < size; ++i) upb_Arena_Free(arenas[i]);
65 }
66
67 class Environment {
68 public:
~Environment()69 ~Environment() {
70 for (auto& atom : arenas_) {
71 auto* a = atom.load(std::memory_order_relaxed);
72 if (a != nullptr) upb_Arena_Free(a);
73 }
74 }
75
RandomNewFree(absl::BitGen & gen)76 void RandomNewFree(absl::BitGen& gen) {
77 auto* old = SwapRandomly(gen, upb_Arena_New());
78 if (old != nullptr) upb_Arena_Free(old);
79 }
80
RandomIncRefCount(absl::BitGen & gen)81 void RandomIncRefCount(absl::BitGen& gen) {
82 auto* a = SwapRandomly(gen, nullptr);
83 if (a != nullptr) {
84 upb_Arena_IncRefFor(a, nullptr);
85 upb_Arena_DecRefFor(a, nullptr);
86 upb_Arena_Free(a);
87 }
88 }
89
RandomFuse(absl::BitGen & gen)90 void RandomFuse(absl::BitGen& gen) {
91 std::array<upb_Arena*, 2> old;
92 for (auto& o : old) {
93 o = SwapRandomly(gen, nullptr);
94 if (o == nullptr) o = upb_Arena_New();
95 }
96
97 EXPECT_TRUE(upb_Arena_Fuse(old[0], old[1]));
98 for (auto& o : old) {
99 o = SwapRandomly(gen, o);
100 if (o != nullptr) upb_Arena_Free(o);
101 }
102 }
103
RandomPoke(absl::BitGen & gen)104 void RandomPoke(absl::BitGen& gen) {
105 switch (absl::Uniform(gen, 0, 2)) {
106 case 0:
107 RandomNewFree(gen);
108 break;
109 case 1:
110 RandomFuse(gen);
111 break;
112 default:
113 break;
114 }
115 }
116
117 private:
SwapRandomly(absl::BitGen & gen,upb_Arena * a)118 upb_Arena* SwapRandomly(absl::BitGen& gen, upb_Arena* a) {
119 return arenas_[absl::Uniform<size_t>(gen, 0, arenas_.size())].exchange(
120 a, std::memory_order_acq_rel);
121 }
122
123 std::array<std::atomic<upb_Arena*>, 100> arenas_ = {};
124 };
125
TEST(ArenaTest,FuzzSingleThreaded)126 TEST(ArenaTest, FuzzSingleThreaded) {
127 Environment env;
128
129 absl::BitGen gen;
130 auto end = absl::Now() + absl::Seconds(0.5);
131 while (absl::Now() < end) {
132 env.RandomPoke(gen);
133 }
134 }
135
TEST(ArenaTest,Contains)136 TEST(ArenaTest, Contains) {
137 upb_Arena* arena1 = upb_Arena_New();
138 upb_Arena* arena2 = upb_Arena_New();
139 void* ptr1a = upb_Arena_Malloc(arena1, 8);
140 void* ptr2a = upb_Arena_Malloc(arena2, 8);
141
142 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
143 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
144 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
145 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
146
147 void* ptr1b = upb_Arena_Malloc(arena1, 1000000);
148 void* ptr2b = upb_Arena_Malloc(arena2, 1000000);
149
150 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
151 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
152 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
153 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
154 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
155 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
156 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
157 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));
158
159 upb_Arena_Fuse(arena1, arena2);
160
161 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
162 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
163 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
164 EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
165 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
166 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
167 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
168 EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));
169
170 upb_Arena_Free(arena1);
171 upb_Arena_Free(arena2);
172 }
173
TEST(ArenaTest,LargeAlloc)174 TEST(ArenaTest, LargeAlloc) {
175 // Tests an allocation larger than the max block size.
176 upb_Arena* arena = upb_Arena_New();
177 size_t size = 100000;
178 char* mem = static_cast<char*>(upb_Arena_Malloc(arena, size));
179 EXPECT_NE(mem, nullptr);
180 for (size_t i = 0; i < size; ++i) {
181 mem[i] = static_cast<char>(i);
182 }
183 for (size_t i = 0; i < size; ++i) {
184 EXPECT_EQ(mem[i], static_cast<char>(i));
185 }
186 upb_Arena_Free(arena);
187 }
188
TEST(ArenaTest,MaxBlockSize)189 TEST(ArenaTest, MaxBlockSize) {
190 upb_Arena* arena = upb_Arena_New();
191 // Perform 600 1k allocations (600k total) and ensure that the amount of
192 // memory allocated does not exceed 700k.
193 for (int i = 0; i < 600; ++i) {
194 upb_Arena_Malloc(arena, 1024);
195 }
196 EXPECT_LE(upb_Arena_SpaceAllocated(arena, nullptr), 700 * 1024);
197 upb_Arena_Free(arena);
198 }
199
200 #ifdef UPB_USE_C11_ATOMICS
201
TEST(ArenaTest,FuzzFuseFreeRace)202 TEST(ArenaTest, FuzzFuseFreeRace) {
203 Environment env;
204
205 absl::Notification done;
206 std::vector<std::thread> threads;
207 for (int i = 0; i < 10; ++i) {
208 threads.emplace_back([&]() {
209 absl::BitGen gen;
210 while (!done.HasBeenNotified()) {
211 env.RandomNewFree(gen);
212 }
213 });
214 }
215
216 absl::BitGen gen;
217 auto end = absl::Now() + absl::Seconds(2);
218 while (absl::Now() < end) {
219 env.RandomFuse(gen);
220 }
221 done.Notify();
222 for (auto& t : threads) t.join();
223 }
224
TEST(ArenaTest,FuzzFuseFuseRace)225 TEST(ArenaTest, FuzzFuseFuseRace) {
226 Environment env;
227
228 absl::Notification done;
229 std::vector<std::thread> threads;
230 for (int i = 0; i < 10; ++i) {
231 threads.emplace_back([&]() {
232 absl::BitGen gen;
233 while (!done.HasBeenNotified()) {
234 env.RandomFuse(gen);
235 }
236 });
237 }
238
239 absl::BitGen gen;
240 auto end = absl::Now() + absl::Seconds(2);
241 while (absl::Now() < end) {
242 env.RandomFuse(gen);
243 }
244 done.Notify();
245 for (auto& t : threads) t.join();
246 }
247
TEST(ArenaTest,ArenaIncRef)248 TEST(ArenaTest, ArenaIncRef) {
249 upb_Arena* arena1 = upb_Arena_New();
250 EXPECT_EQ(upb_Arena_DebugRefCount(arena1), 1);
251 upb_Arena_IncRefFor(arena1, nullptr);
252 EXPECT_EQ(upb_Arena_DebugRefCount(arena1), 2);
253 upb_Arena_DecRefFor(arena1, nullptr);
254 EXPECT_EQ(upb_Arena_DebugRefCount(arena1), 1);
255 upb_Arena_Free(arena1);
256 }
257
TEST(ArenaTest,FuzzFuseIncRefCountRace)258 TEST(ArenaTest, FuzzFuseIncRefCountRace) {
259 Environment env;
260
261 absl::Notification done;
262 std::vector<std::thread> threads;
263 for (int i = 0; i < 10; ++i) {
264 threads.emplace_back([&]() {
265 absl::BitGen gen;
266 while (!done.HasBeenNotified()) {
267 env.RandomNewFree(gen);
268 }
269 });
270 }
271
272 absl::BitGen gen;
273 auto end = absl::Now() + absl::Seconds(2);
274 while (absl::Now() < end) {
275 env.RandomFuse(gen);
276 env.RandomIncRefCount(gen);
277 }
278 done.Notify();
279 for (auto& t : threads) t.join();
280 }
281
TEST(ArenaTest,IncRefCountShouldFailForInitialBlock)282 TEST(ArenaTest, IncRefCountShouldFailForInitialBlock) {
283 char buf1[1024];
284 upb_Arena* arena = upb_Arena_Init(buf1, 1024, &upb_alloc_global);
285 EXPECT_FALSE(upb_Arena_IncRefFor(arena, nullptr));
286 }
287
288 #endif
289
290 } // namespace
291