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