2 * Copyright 2017 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include <folly/io/RecordIO.h>
19 #include <sys/types.h>
23 #include <glog/logging.h>
25 #include <folly/Conv.h>
26 #include <folly/FBString.h>
27 #include <folly/Random.h>
28 #include <folly/experimental/TestUtil.h>
29 #include <folly/io/IOBufQueue.h>
30 #include <folly/portability/GFlags.h>
31 #include <folly/portability/GTest.h>
32 #include <folly/portability/Unistd.h>
34 DEFINE_int32(random_seed, folly::randomNumberSeed(), "random seed");
41 StringPiece sp(ByteRange br) { return StringPiece(br); }
44 std::unique_ptr<IOBuf> iobufs(std::initializer_list<T> ranges) {
46 for (auto& range : ranges) {
48 queue.append(IOBuf::wrapBuffer(r.data(), r.size()));
55 TEST(RecordIOTest, Simple) {
58 RecordIOWriter writer(File(file.fd()));
59 writer.write(iobufs({"hello ", "world"}));
60 writer.write(iobufs({"goodbye"}));
63 RecordIOReader reader(File(file.fd()));
64 auto it = reader.begin();
65 ASSERT_FALSE(it == reader.end());
66 EXPECT_EQ("hello world", sp((it++)->first));
67 ASSERT_FALSE(it == reader.end());
68 EXPECT_EQ("goodbye", sp((it++)->first));
69 EXPECT_TRUE(it == reader.end());
72 RecordIOWriter writer(File(file.fd()));
73 writer.write(iobufs({"meow"}));
74 writer.write(iobufs({"woof"}));
77 RecordIOReader reader(File(file.fd()));
78 auto it = reader.begin();
79 ASSERT_FALSE(it == reader.end());
80 EXPECT_EQ("hello world", sp((it++)->first));
81 ASSERT_FALSE(it == reader.end());
82 EXPECT_EQ("goodbye", sp((it++)->first));
83 ASSERT_FALSE(it == reader.end());
84 EXPECT_EQ("meow", sp((it++)->first));
85 ASSERT_FALSE(it == reader.end());
86 EXPECT_EQ("woof", sp((it++)->first));
87 EXPECT_TRUE(it == reader.end());
91 TEST(RecordIOTest, SmallRecords) {
92 constexpr size_t kSize = 10;
94 memset(tmp, 'x', kSize);
97 RecordIOWriter writer(File(file.fd()));
98 for (size_t i = 0; i < kSize; ++i) { // record of size 0 should be ignored
99 writer.write(IOBuf::wrapBuffer(tmp, i));
103 RecordIOReader reader(File(file.fd()));
104 auto it = reader.begin();
105 for (size_t i = 1; i < kSize; ++i) {
106 ASSERT_FALSE(it == reader.end());
107 EXPECT_EQ(StringPiece(tmp, i), sp((it++)->first));
109 EXPECT_TRUE(it == reader.end());
113 TEST(RecordIOTest, MultipleFileIds) {
116 RecordIOWriter writer(File(file.fd()), 1);
117 writer.write(iobufs({"hello"}));
120 RecordIOWriter writer(File(file.fd()), 2);
121 writer.write(iobufs({"world"}));
124 RecordIOWriter writer(File(file.fd()), 1);
125 writer.write(iobufs({"goodbye"}));
128 RecordIOReader reader(File(file.fd()), 0); // return all
129 auto it = reader.begin();
130 ASSERT_FALSE(it == reader.end());
131 EXPECT_EQ("hello", sp((it++)->first));
132 ASSERT_FALSE(it == reader.end());
133 EXPECT_EQ("world", sp((it++)->first));
134 ASSERT_FALSE(it == reader.end());
135 EXPECT_EQ("goodbye", sp((it++)->first));
136 EXPECT_TRUE(it == reader.end());
139 RecordIOReader reader(File(file.fd()), 1);
140 auto it = reader.begin();
141 ASSERT_FALSE(it == reader.end());
142 EXPECT_EQ("hello", sp((it++)->first));
143 ASSERT_FALSE(it == reader.end());
144 EXPECT_EQ("goodbye", sp((it++)->first));
145 EXPECT_TRUE(it == reader.end());
148 RecordIOReader reader(File(file.fd()), 2);
149 auto it = reader.begin();
150 ASSERT_FALSE(it == reader.end());
151 EXPECT_EQ("world", sp((it++)->first));
152 EXPECT_TRUE(it == reader.end());
155 RecordIOReader reader(File(file.fd()), 3);
156 auto it = reader.begin();
157 EXPECT_TRUE(it == reader.end());
161 TEST(RecordIOTest, ExtraMagic) {
164 RecordIOWriter writer(File(file.fd()));
165 writer.write(iobufs({"hello"}));
167 uint8_t buf[recordio_helpers::headerSize() + 5];
168 EXPECT_EQ(0, lseek(file.fd(), 0, SEEK_SET));
169 EXPECT_EQ(sizeof(buf), read(file.fd(), buf, sizeof(buf)));
170 // Append an extra magic
171 const uint32_t magic = recordio_helpers::detail::Header::kMagic;
172 EXPECT_EQ(sizeof(magic), write(file.fd(), &magic, sizeof(magic)));
173 // and an extra record
174 EXPECT_EQ(sizeof(buf), write(file.fd(), buf, sizeof(buf)));
176 RecordIOReader reader(File(file.fd()));
177 auto it = reader.begin();
178 ASSERT_FALSE(it == reader.end());
179 EXPECT_EQ("hello", sp((it++)->first));
180 ASSERT_FALSE(it == reader.end());
181 EXPECT_EQ("hello", sp((it++)->first));
182 EXPECT_TRUE(it == reader.end());
187 void corrupt(int fd, off_t pos) {
189 EXPECT_EQ(1, pread(fd, &val, 1, pos));
191 EXPECT_EQ(1, pwrite(fd, &val, 1, pos));
195 TEST(RecordIOTest, Randomized) {
196 SCOPED_TRACE(to<std::string>("Random seed is ", FLAGS_random_seed));
197 std::mt19937 rnd(FLAGS_random_seed);
200 std::uniform_int_distribution<uint32_t>(30, 300)(rnd);
202 std::uniform_int_distribution<uint32_t> recordSizeDist(1, 3 << 16);
203 std::uniform_int_distribution<uint32_t> charDist(0, 255);
204 std::uniform_int_distribution<uint32_t> junkDist(0, 1 << 20);
205 // corrupt 1/5 of all records
206 std::uniform_int_distribution<uint32_t> corruptDist(0, 4);
208 std::vector<std::pair<fbstring, off_t>> records;
209 std::vector<off_t> corruptPositions;
210 records.reserve(recordCount);
214 // Recreate the writer multiple times so we test that we create a
216 for (size_t i = 0; i < 3; ++i) {
217 RecordIOWriter writer(File(file.fd()));
218 for (size_t j = 0; j < recordCount; ++j) {
219 off_t beginPos = writer.filePos();
221 size_t recordSize = recordSizeDist(rnd);
222 record.reserve(recordSize);
223 for (size_t k = 0; k < recordSize; ++k) {
224 record.push_back(charDist(rnd));
226 writer.write(iobufs({record}));
228 bool corrupt = (corruptDist(rnd) == 0);
230 // Corrupt one random byte in the record (including header)
231 std::uniform_int_distribution<uint32_t> corruptByteDist(
232 0, recordSize + recordio_helpers::headerSize() - 1);
233 off_t corruptRel = corruptByteDist(rnd);
234 VLOG(1) << "n=" << records.size() << " bpos=" << beginPos
235 << " rsize=" << record.size()
236 << " corrupt rel=" << corruptRel
237 << " abs=" << beginPos + corruptRel;
238 corruptPositions.push_back(beginPos + corruptRel);
240 VLOG(2) << "n=" << records.size() << " bpos=" << beginPos
241 << " rsize=" << record.size()
243 records.emplace_back(std::move(record), beginPos);
246 VLOG(1) << "n=" << records.size() << " close abs=" << writer.filePos();
249 for (auto& pos : corruptPositions) {
250 corrupt(file.fd(), pos);
255 RecordIOReader reader(File(file.fd()));
256 for (auto& r : reader) {
258 ASSERT_LT(i, records.size());
259 EXPECT_EQ(records[i].first, sp(r.first));
260 EXPECT_EQ(records[i].second, r.second);
263 EXPECT_EQ(records.size(), i);
270 int main(int argc, char *argv[]) {
271 testing::InitGoogleTest(&argc, argv);
272 gflags::ParseCommandLineFlags(&argc, &argv, true);
273 return RUN_ALL_TESTS();