Use the GTest portability headers
[folly.git] / folly / io / test / RecordIOTest.cpp
1 /*
2  * Copyright 2016 Facebook, Inc.
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/RecordIO.h>
18
19 #include <sys/types.h>
20
21 #include <random>
22
23 #include <glog/logging.h>
24
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>
33
34 DEFINE_int32(random_seed, folly::randomNumberSeed(), "random seed");
35
36 namespace folly { namespace test {
37
38 namespace {
39 // shortcut
40 StringPiece sp(ByteRange br) { return StringPiece(br); }
41
42 template <class T>
43 std::unique_ptr<IOBuf> iobufs(std::initializer_list<T> ranges) {
44   IOBufQueue queue;
45   for (auto& range : ranges) {
46     StringPiece r(range);
47     queue.append(IOBuf::wrapBuffer(r.data(), r.size()));
48   }
49   return queue.move();
50 }
51
52 }  // namespace
53
54 TEST(RecordIOTest, Simple) {
55   TemporaryFile file;
56   {
57     RecordIOWriter writer(File(file.fd()));
58     writer.write(iobufs({"hello ", "world"}));
59     writer.write(iobufs({"goodbye"}));
60   }
61   {
62     RecordIOReader reader(File(file.fd()));
63     auto it = reader.begin();
64     ASSERT_FALSE(it == reader.end());
65     EXPECT_EQ("hello world", sp((it++)->first));
66     ASSERT_FALSE(it == reader.end());
67     EXPECT_EQ("goodbye", sp((it++)->first));
68     EXPECT_TRUE(it == reader.end());
69   }
70   {
71     RecordIOWriter writer(File(file.fd()));
72     writer.write(iobufs({"meow"}));
73     writer.write(iobufs({"woof"}));
74   }
75   {
76     RecordIOReader reader(File(file.fd()));
77     auto it = reader.begin();
78     ASSERT_FALSE(it == reader.end());
79     EXPECT_EQ("hello world", sp((it++)->first));
80     ASSERT_FALSE(it == reader.end());
81     EXPECT_EQ("goodbye", sp((it++)->first));
82     ASSERT_FALSE(it == reader.end());
83     EXPECT_EQ("meow", sp((it++)->first));
84     ASSERT_FALSE(it == reader.end());
85     EXPECT_EQ("woof", sp((it++)->first));
86     EXPECT_TRUE(it == reader.end());
87   }
88 }
89
90 TEST(RecordIOTest, SmallRecords) {
91   constexpr size_t kSize = 10;
92   char tmp[kSize];
93   memset(tmp, 'x', kSize);
94   TemporaryFile file;
95   {
96     RecordIOWriter writer(File(file.fd()));
97     for (size_t i = 0; i < kSize; ++i) {  // record of size 0 should be ignored
98       writer.write(IOBuf::wrapBuffer(tmp, i));
99     }
100   }
101   {
102     RecordIOReader reader(File(file.fd()));
103     auto it = reader.begin();
104     for (size_t i = 1; i < kSize; ++i) {
105       ASSERT_FALSE(it == reader.end());
106       EXPECT_EQ(StringPiece(tmp, i), sp((it++)->first));
107     }
108     EXPECT_TRUE(it == reader.end());
109   }
110 }
111
112 TEST(RecordIOTest, MultipleFileIds) {
113   TemporaryFile file;
114   {
115     RecordIOWriter writer(File(file.fd()), 1);
116     writer.write(iobufs({"hello"}));
117   }
118   {
119     RecordIOWriter writer(File(file.fd()), 2);
120     writer.write(iobufs({"world"}));
121   }
122   {
123     RecordIOWriter writer(File(file.fd()), 1);
124     writer.write(iobufs({"goodbye"}));
125   }
126   {
127     RecordIOReader reader(File(file.fd()), 0);  // return all
128     auto it = reader.begin();
129     ASSERT_FALSE(it == reader.end());
130     EXPECT_EQ("hello", sp((it++)->first));
131     ASSERT_FALSE(it == reader.end());
132     EXPECT_EQ("world", sp((it++)->first));
133     ASSERT_FALSE(it == reader.end());
134     EXPECT_EQ("goodbye", sp((it++)->first));
135     EXPECT_TRUE(it == reader.end());
136   }
137   {
138     RecordIOReader reader(File(file.fd()), 1);
139     auto it = reader.begin();
140     ASSERT_FALSE(it == reader.end());
141     EXPECT_EQ("hello", sp((it++)->first));
142     ASSERT_FALSE(it == reader.end());
143     EXPECT_EQ("goodbye", sp((it++)->first));
144     EXPECT_TRUE(it == reader.end());
145   }
146   {
147     RecordIOReader reader(File(file.fd()), 2);
148     auto it = reader.begin();
149     ASSERT_FALSE(it == reader.end());
150     EXPECT_EQ("world", sp((it++)->first));
151     EXPECT_TRUE(it == reader.end());
152   }
153   {
154     RecordIOReader reader(File(file.fd()), 3);
155     auto it = reader.begin();
156     EXPECT_TRUE(it == reader.end());
157   }
158 }
159
160 TEST(RecordIOTest, ExtraMagic) {
161   TemporaryFile file;
162   {
163     RecordIOWriter writer(File(file.fd()));
164     writer.write(iobufs({"hello"}));
165   }
166   uint8_t buf[recordio_helpers::headerSize() + 5];
167   EXPECT_EQ(0, lseek(file.fd(), 0, SEEK_SET));
168   EXPECT_EQ(sizeof(buf), read(file.fd(), buf, sizeof(buf)));
169   // Append an extra magic
170   const uint32_t magic = recordio_helpers::detail::Header::kMagic;
171   EXPECT_EQ(sizeof(magic), write(file.fd(), &magic, sizeof(magic)));
172   // and an extra record
173   EXPECT_EQ(sizeof(buf), write(file.fd(), buf, sizeof(buf)));
174   {
175     RecordIOReader reader(File(file.fd()));
176     auto it = reader.begin();
177     ASSERT_FALSE(it == reader.end());
178     EXPECT_EQ("hello", sp((it++)->first));
179     ASSERT_FALSE(it == reader.end());
180     EXPECT_EQ("hello", sp((it++)->first));
181     EXPECT_TRUE(it == reader.end());
182   }
183 }
184
185 namespace {
186 void corrupt(int fd, off_t pos) {
187   uint8_t val = 0;
188   EXPECT_EQ(1, pread(fd, &val, 1, pos));
189   ++val;
190   EXPECT_EQ(1, pwrite(fd, &val, 1, pos));
191 }
192 }  // namespace
193
194 TEST(RecordIOTest, Randomized) {
195   SCOPED_TRACE(to<std::string>("Random seed is ", FLAGS_random_seed));
196   std::mt19937 rnd(FLAGS_random_seed);
197
198   size_t recordCount =
199     std::uniform_int_distribution<uint32_t>(30, 300)(rnd);
200
201   std::uniform_int_distribution<uint32_t> recordSizeDist(1, 3 << 16);
202   std::uniform_int_distribution<uint32_t> charDist(0, 255);
203   std::uniform_int_distribution<uint32_t> junkDist(0, 1 << 20);
204   // corrupt 1/5 of all records
205   std::uniform_int_distribution<uint32_t> corruptDist(0, 4);
206
207   std::vector<std::pair<fbstring, off_t>> records;
208   std::vector<off_t> corruptPositions;
209   records.reserve(recordCount);
210   TemporaryFile file;
211
212   fbstring record;
213   // Recreate the writer multiple times so we test that we create a
214   // continuous stream
215   for (size_t i = 0; i < 3; ++i) {
216     RecordIOWriter writer(File(file.fd()));
217     for (size_t j = 0; j < recordCount; ++j) {
218       off_t beginPos = writer.filePos();
219       record.clear();
220       size_t recordSize = recordSizeDist(rnd);
221       record.reserve(recordSize);
222       for (size_t k = 0; k < recordSize; ++k) {
223         record.push_back(charDist(rnd));
224       }
225       writer.write(iobufs({record}));
226
227       bool corrupt = (corruptDist(rnd) == 0);
228       if (corrupt) {
229         // Corrupt one random byte in the record (including header)
230         std::uniform_int_distribution<uint32_t> corruptByteDist(
231             0, recordSize + recordio_helpers::headerSize() - 1);
232         off_t corruptRel = corruptByteDist(rnd);
233         VLOG(1) << "n=" << records.size() << " bpos=" << beginPos
234                 << " rsize=" << record.size()
235                 << " corrupt rel=" << corruptRel
236                 << " abs=" << beginPos + corruptRel;
237         corruptPositions.push_back(beginPos + corruptRel);
238       } else {
239         VLOG(2) << "n=" << records.size() << " bpos=" << beginPos
240                 << " rsize=" << record.size()
241                 << " good";
242         records.emplace_back(std::move(record), beginPos);
243       }
244     }
245     VLOG(1) << "n=" << records.size() << " close abs=" << writer.filePos();
246   }
247
248   for (auto& pos : corruptPositions) {
249     corrupt(file.fd(), pos);
250   }
251
252   {
253     size_t i = 0;
254     RecordIOReader reader(File(file.fd()));
255     for (auto& r : reader) {
256       SCOPED_TRACE(i);
257       ASSERT_LT(i, records.size());
258       EXPECT_EQ(records[i].first, sp(r.first));
259       EXPECT_EQ(records[i].second, r.second);
260       ++i;
261     }
262     EXPECT_EQ(records.size(), i);
263   }
264 }
265
266 }}  // namespaces
267
268 int main(int argc, char *argv[]) {
269   testing::InitGoogleTest(&argc, argv);
270   gflags::ParseCommandLineFlags(&argc, &argv, true);
271   return RUN_ALL_TESTS();
272 }