folly/io/test/RecordIOTest.cpp

   1 /*
   2  * Copyright 2014 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 #include <unistd.h>
  21
  22 #include <random>
  23
  24 #include <gflags/gflags.h>
  25 #include <glog/logging.h>
  26 #include <gtest/gtest.h>
  27
  28 #include <folly/Conv.h>
  29 #include <folly/FBString.h>
  30 #include <folly/Random.h>
  31 #include <folly/experimental/TestUtil.h>
  32 #include <folly/io/IOBufQueue.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 (int 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 (int 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 }