folly/io/test/RecordIOTest.cpp

   1 /*
   2  * Copyright 2013 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 ByteRange br(StringPiece sp) { return ByteRange(sp); }
  41 StringPiece sp(ByteRange br) { return StringPiece(br); }
  42
  43 template <class T>
  44 std::unique_ptr<IOBuf> iobufs(std::initializer_list<T> ranges) {
  45   IOBufQueue queue;
  46   for (auto& range : ranges) {
  47     StringPiece r(range);
  48     queue.append(IOBuf::wrapBuffer(r.data(), r.size()));
  49   }
  50   return queue.move();
  51 }
  52
  53 }  // namespace
  54
  55 TEST(RecordIOTest, Simple) {
  56   TemporaryFile file;
  57   {
  58     RecordIOWriter writer(file.fd());
  59     writer.write(iobufs({"hello ", "world"}));
  60     writer.write(iobufs({"goodbye"}));
  61   }
  62   {
  63     RecordIOReader reader(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());
  70   }
  71   {
  72     RecordIOWriter writer(file.fd());
  73     writer.write(iobufs({"meow"}));
  74     writer.write(iobufs({"woof"}));
  75   }
  76   {
  77     RecordIOReader reader(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());
  88   }
  89 }
  90
  91 TEST(RecordIOTest, SmallRecords) {
  92   constexpr size_t kSize = 10;
  93   char tmp[kSize];
  94   memset(tmp, 'x', kSize);
  95   TemporaryFile file;
  96   {
  97     RecordIOWriter writer(file.fd());
  98     for (int i = 0; i < kSize; ++i) {  // record of size 0 should be ignored
  99       writer.write(IOBuf::wrapBuffer(tmp, i));
 100     }
 101   }
 102   {
 103     RecordIOReader reader(file.fd());
 104     auto it = reader.begin();
 105     for (int i = 1; i < kSize; ++i) {
 106       ASSERT_FALSE(it == reader.end());
 107       EXPECT_EQ(StringPiece(tmp, i), sp((it++)->first));
 108     }
 109     EXPECT_TRUE(it == reader.end());
 110   }
 111 }
 112
 113 TEST(RecordIOTest, MultipleFileIds) {
 114   TemporaryFile file;
 115   {
 116     RecordIOWriter writer(file.fd(), 1);
 117     writer.write(iobufs({"hello"}));
 118   }
 119   {
 120     RecordIOWriter writer(file.fd(), 2);
 121     writer.write(iobufs({"world"}));
 122   }
 123   {
 124     RecordIOWriter writer(file.fd(), 1);
 125     writer.write(iobufs({"goodbye"}));
 126   }
 127   {
 128     RecordIOReader reader(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());
 137   }
 138   {
 139     RecordIOReader reader(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());
 146   }
 147   {
 148     RecordIOReader reader(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());
 153   }
 154   {
 155     RecordIOReader reader(file.fd(), 3);
 156     auto it = reader.begin();
 157     EXPECT_TRUE(it == reader.end());
 158   }
 159 }
 160
 161 TEST(RecordIOTest, ExtraMagic) {
 162   TemporaryFile file;
 163   {
 164     RecordIOWriter writer(file.fd());
 165     writer.write(iobufs({"hello"}));
 166   }
 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)));
 175   {
 176     RecordIOReader reader(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());
 183   }
 184 }
 185
 186 namespace {
 187 void corrupt(int fd, off_t pos) {
 188   uint8_t val = 0;
 189   EXPECT_EQ(1, pread(fd, &val, 1, pos));
 190   ++val;
 191   EXPECT_EQ(1, pwrite(fd, &val, 1, pos));
 192 }
 193 }  // namespace
 194
 195 TEST(RecordIOTest, Randomized) {
 196   SCOPED_TRACE(to<std::string>("Random seed is ", FLAGS_random_seed));
 197   std::mt19937 rnd(FLAGS_random_seed);
 198
 199   size_t recordCount =
 200     std::uniform_int_distribution<uint32_t>(30, 300)(rnd);
 201
 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);
 207
 208   std::vector<std::pair<fbstring, off_t>> records;
 209   std::vector<off_t> corruptPositions;
 210   records.reserve(recordCount);
 211   TemporaryFile file;
 212
 213   fbstring record;
 214   // Recreate the writer multiple times so we test that we create a
 215   // continuous stream
 216   for (size_t i = 0; i < 3; ++i) {
 217     RecordIOWriter writer(file.fd());
 218     for (size_t j = 0; j < recordCount; ++j) {
 219       off_t beginPos = writer.filePos();
 220       record.clear();
 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));
 225       }
 226       writer.write(iobufs({record}));
 227
 228       bool corrupt = (corruptDist(rnd) == 0);
 229       if (corrupt) {
 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);
 239       } else {
 240         VLOG(2) << "n=" << records.size() << " bpos=" << beginPos
 241                 << " rsize=" << record.size()
 242                 << " good";
 243         records.emplace_back(std::move(record), beginPos);
 244       }
 245     }
 246     VLOG(1) << "n=" << records.size() << " close abs=" << writer.filePos();
 247   }
 248
 249   for (auto& pos : corruptPositions) {
 250     corrupt(file.fd(), pos);
 251   }
 252
 253   {
 254     size_t i = 0;
 255     RecordIOReader reader(file.fd());
 256     for (auto& r : reader) {
 257       SCOPED_TRACE(i);
 258       ASSERT_LT(i, records.size());
 259       EXPECT_EQ(records[i].first, sp(r.first));
 260       EXPECT_EQ(records[i].second, r.second);
 261       ++i;
 262     }
 263     EXPECT_EQ(records.size(), i);
 264   }
 265 }
 266
 267 }}  // namespaces
 268
 269 int main(int argc, char *argv[]) {
 270   testing::InitGoogleTest(&argc, argv);
 271   google::ParseCommandLineFlags(&argc, &argv, true);
 272   return RUN_ALL_TESTS();
 273 }
 274