1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <folly/io/IOBuf.h>
18
19 #include <folly/portability/GTest.h>
20
21 std::atomic<size_t> gIOBufAlloc{};
22
io_buf_alloc_cb(void *,size_t size)23 void io_buf_alloc_cb(void* /*ptr*/, size_t size) noexcept {
24 gIOBufAlloc += size;
25 }
26
io_buf_free_cb(void *,size_t size)27 void io_buf_free_cb(void* /*ptr*/, size_t size) noexcept {
28 gIOBufAlloc -= size;
29 }
30
31 namespace {
32 struct IOBufTestInfo {
33 std::unique_ptr<folly::IOBuf> ioBuf;
34 size_t delta{0};
35
reset__anon427e56710111::IOBufTestInfo36 void reset() { ioBuf.reset(); }
37 };
38 } // namespace
39
TEST(IOBufCB,AllocSizes)40 TEST(IOBufCB, AllocSizes) {
41 std::vector<IOBufTestInfo> testInfoVec;
42 size_t initialSize = gIOBufAlloc.load();
43 size_t prevSize = initialSize;
44
45 auto allocAndAppend = [&](size_t size) {
46 auto buf = folly::IOBuf::create(size);
47 size_t newSize = gIOBufAlloc.load();
48 CHECK_GT(newSize, prevSize);
49 size_t delta = newSize - prevSize;
50 if (delta) {
51 CHECK_GE(delta, size);
52 }
53 prevSize = newSize;
54
55 IOBufTestInfo info;
56 info.ioBuf = std::move(buf);
57 info.delta = delta;
58
59 testInfoVec.emplace_back(std::move(info));
60 };
61 // try with smaller allocations
62 for (size_t i = 0; i < 1234; i++) {
63 allocAndAppend(i);
64 }
65
66 // try with large allocations that will require an external buffer
67 for (size_t i = 1; i < 100; i++) {
68 allocAndAppend(i * 1024);
69 }
70
71 // deallocate
72 for (size_t i = 0; i < testInfoVec.size(); i++) {
73 testInfoVec[i].reset();
74 }
75
76 CHECK_EQ(initialSize, gIOBufAlloc.load());
77 }
78
TEST(IOBufCB,TakeOwnership)79 TEST(IOBufCB, TakeOwnership) {
80 size_t initialSize = gIOBufAlloc.load();
81 static constexpr size_t kAllocSize = 1024;
82
83 {
84 auto buf = folly::IOBuf::takeOwnership(
85 folly::IOBuf::SIZED_FREE, malloc(kAllocSize), kAllocSize, 0, 0);
86 size_t newSize = gIOBufAlloc.load();
87 CHECK_GE(newSize, initialSize + kAllocSize);
88 }
89
90 CHECK_EQ(initialSize, gIOBufAlloc.load());
91
92 {
93 folly::IOBuf buf(
94 folly::IOBuf::TAKE_OWNERSHIP,
95 folly::IOBuf::SIZED_FREE,
96 malloc(kAllocSize),
97 kAllocSize,
98 0,
99 0);
100 size_t newSize = gIOBufAlloc.load();
101 CHECK_GE(newSize, initialSize + kAllocSize);
102 }
103
104 CHECK_EQ(initialSize, gIOBufAlloc.load());
105 }
106
TEST(IOBufCB,MoveToFbString)107 TEST(IOBufCB, MoveToFbString) {
108 size_t initialSize = gIOBufAlloc.load();
109 static constexpr size_t kAllocSize = 1024;
110
111 LOG(INFO) << gIOBufAlloc.load();
112
113 {
114 auto buf = folly::IOBuf::takeOwnership(
115 folly::IOBuf::SIZED_FREE, malloc(kAllocSize), kAllocSize, 0, 0);
116 LOG(INFO) << gIOBufAlloc.load();
117 size_t newSize = gIOBufAlloc.load();
118 CHECK_GE(newSize, initialSize + kAllocSize);
119
120 auto str = buf->moveToFbString();
121 LOG(INFO) << gIOBufAlloc.load();
122 CHECK_LT(gIOBufAlloc.load(), newSize);
123 }
124 LOG(INFO) << gIOBufAlloc.load();
125
126 CHECK_EQ(initialSize, gIOBufAlloc.load());
127 }
128