folly/io/test/CompressionTest.cpp

   1 /*
   2  * Copyright 2015 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/Compression.h>
  18
  19 #include <random>
  20 #include <thread>
  21 #include <unordered_map>
  22
  23 #include <boost/noncopyable.hpp>
  24 #include <glog/logging.h>
  25 #include <gtest/gtest.h>
  26
  27 #include <folly/Benchmark.h>
  28 #include <folly/Hash.h>
  29 #include <folly/Random.h>
  30 #include <folly/Varint.h>
  31 #include <folly/io/IOBufQueue.h>
  32
  33 namespace folly { namespace io { namespace test {
  34
  35 class DataHolder : private boost::noncopyable {
  36  public:
  37   uint64_t hash(size_t size) const;
  38   ByteRange data(size_t size) const;
  39
  40  protected:
  41   explicit DataHolder(size_t sizeLog2);
  42   const size_t size_;
  43   std::unique_ptr<uint8_t[]> data_;
  44   mutable std::unordered_map<uint64_t, uint64_t> hashCache_;
  45 };
  46
  47 DataHolder::DataHolder(size_t sizeLog2)
  48   : size_(size_t(1) << sizeLog2),
  49     data_(new uint8_t[size_]) {
  50 }
  51
  52 uint64_t DataHolder::hash(size_t size) const {
  53   CHECK_LE(size, size_);
  54   auto p = hashCache_.find(size);
  55   if (p != hashCache_.end()) {
  56     return p->second;
  57   }
  58
  59   uint64_t h = folly::hash::fnv64_buf(data_.get(), size);
  60   hashCache_[size] = h;
  61   return h;
  62 }
  63
  64 ByteRange DataHolder::data(size_t size) const {
  65   CHECK_LE(size, size_);
  66   return ByteRange(data_.get(), size);
  67 }
  68
  69 uint64_t hashIOBuf(const IOBuf* buf) {
  70   uint64_t h = folly::hash::FNV_64_HASH_START;
  71   for (auto& range : *buf) {
  72     h = folly::hash::fnv64_buf(range.data(), range.size(), h);
  73   }
  74   return h;
  75 }
  76
  77 class RandomDataHolder : public DataHolder {
  78  public:
  79   explicit RandomDataHolder(size_t sizeLog2);
  80 };
  81
  82 RandomDataHolder::RandomDataHolder(size_t sizeLog2)
  83   : DataHolder(sizeLog2) {
  84   constexpr size_t numThreadsLog2 = 3;
  85   constexpr size_t numThreads = size_t(1) << numThreadsLog2;
  86
  87   uint32_t seed = randomNumberSeed();
  88
  89   std::vector<std::thread> threads;
  90   threads.reserve(numThreads);
  91   for (size_t t = 0; t < numThreads; ++t) {
  92     threads.emplace_back(
  93         [this, seed, t, numThreadsLog2, sizeLog2] () {
  94           std::mt19937 rng(seed + t);
  95           size_t countLog2 = size_t(1) << (sizeLog2 - numThreadsLog2);
  96           size_t start = size_t(t) << countLog2;
  97           for (size_t i = 0; i < countLog2; ++i) {
  98             this->data_[start + i] = rng();
  99           }
 100         });
 101   }
 102
 103   for (auto& t : threads) {
 104     t.join();
 105   }
 106 }
 107
 108 class ConstantDataHolder : public DataHolder {
 109  public:
 110   explicit ConstantDataHolder(size_t sizeLog2);
 111 };
 112
 113 ConstantDataHolder::ConstantDataHolder(size_t sizeLog2)
 114   : DataHolder(sizeLog2) {
 115   memset(data_.get(), 'a', size_);
 116 }
 117
 118 constexpr size_t dataSizeLog2 = 27;  // 128MiB
 119 RandomDataHolder randomDataHolder(dataSizeLog2);
 120 ConstantDataHolder constantDataHolder(dataSizeLog2);
 121
 122 TEST(CompressionTestNeedsUncompressedLength, Simple) {
 123   EXPECT_FALSE(getCodec(CodecType::NO_COMPRESSION)->needsUncompressedLength());
 124   EXPECT_TRUE(getCodec(CodecType::LZ4)->needsUncompressedLength());
 125   EXPECT_FALSE(getCodec(CodecType::SNAPPY)->needsUncompressedLength());
 126   EXPECT_FALSE(getCodec(CodecType::ZLIB)->needsUncompressedLength());
 127   EXPECT_FALSE(getCodec(CodecType::LZ4_VARINT_SIZE)->needsUncompressedLength());
 128   EXPECT_TRUE(getCodec(CodecType::LZMA2)->needsUncompressedLength());
 129   EXPECT_FALSE(getCodec(CodecType::LZMA2_VARINT_SIZE)
 130     ->needsUncompressedLength());
 131 }
 132
 133 class CompressionTest
 134     : public testing::TestWithParam<std::tr1::tuple<int, CodecType>> {
 135   protected:
 136    void SetUp() override {
 137      auto tup = GetParam();
 138      uncompressedLength_ = uint64_t(1) << std::tr1::get<0>(tup);
 139      codec_ = getCodec(std::tr1::get<1>(tup));
 140    }
 141
 142    void runSimpleTest(const DataHolder& dh);
 143
 144    uint64_t uncompressedLength_;
 145    std::unique_ptr<Codec> codec_;
 146 };
 147
 148 void CompressionTest::runSimpleTest(const DataHolder& dh) {
 149   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength_));
 150   auto compressed = codec_->compress(original.get());
 151   if (!codec_->needsUncompressedLength()) {
 152     auto uncompressed = codec_->uncompress(compressed.get());
 153
 154     EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 155     EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 156   }
 157   {
 158     auto uncompressed = codec_->uncompress(compressed.get(),
 159                                            uncompressedLength_);
 160     EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 161     EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 162   }
 163 }
 164
 165 TEST_P(CompressionTest, RandomData) {
 166   runSimpleTest(randomDataHolder);
 167 }
 168
 169 TEST_P(CompressionTest, ConstantData) {
 170   runSimpleTest(constantDataHolder);
 171 }
 172
 173 INSTANTIATE_TEST_CASE_P(
 174     CompressionTest,
 175     CompressionTest,
 176     testing::Combine(testing::Values(0, 1, 12, 22, 25, 27),
 177                      testing::Values(CodecType::NO_COMPRESSION,
 178                                      CodecType::LZ4,
 179                                      CodecType::SNAPPY,
 180                                      CodecType::ZLIB,
 181                                      CodecType::LZ4_VARINT_SIZE,
 182                                      CodecType::LZMA2,
 183                                      CodecType::LZMA2_VARINT_SIZE)));
 184
 185 class CompressionVarintTest
 186     : public testing::TestWithParam<std::tr1::tuple<int, CodecType>> {
 187  protected:
 188   void SetUp() override {
 189     auto tup = GetParam();
 190     uncompressedLength_ = uint64_t(1) << std::tr1::get<0>(tup);
 191     codec_ = getCodec(std::tr1::get<1>(tup));
 192   }
 193
 194   void runSimpleTest(const DataHolder& dh);
 195
 196   uint64_t uncompressedLength_;
 197   std::unique_ptr<Codec> codec_;
 198 };
 199
 200 inline uint64_t oneBasedMsbPos(uint64_t number) {
 201   uint64_t pos = 0;
 202   for (; number > 0; ++pos, number >>= 1) {
 203   }
 204   return pos;
 205 }
 206
 207 void CompressionVarintTest::runSimpleTest(const DataHolder& dh) {
 208   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength_));
 209   auto compressed = codec_->compress(original.get());
 210   auto breakPoint =
 211       1UL +
 212       Random::rand64(std::max(9UL, oneBasedMsbPos(uncompressedLength_)) / 9UL);
 213   auto tinyBuf = IOBuf::copyBuffer(compressed->data(),
 214                                    std::min(compressed->length(), breakPoint));
 215   compressed->trimStart(breakPoint);
 216   tinyBuf->prependChain(std::move(compressed));
 217   compressed = std::move(tinyBuf);
 218
 219   auto uncompressed = codec_->uncompress(compressed.get());
 220
 221   EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 222   EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 223 }
 224
 225 TEST_P(CompressionVarintTest, RandomData) { runSimpleTest(randomDataHolder); }
 226
 227 TEST_P(CompressionVarintTest, ConstantData) {
 228   runSimpleTest(constantDataHolder);
 229 }
 230
 231 INSTANTIATE_TEST_CASE_P(
 232     CompressionVarintTest,
 233     CompressionVarintTest,
 234     testing::Combine(testing::Values(0, 1, 12, 22, 25, 27),
 235                      testing::Values(CodecType::LZ4_VARINT_SIZE,
 236                                      CodecType::LZMA2_VARINT_SIZE)));
 237
 238 class CompressionCorruptionTest : public testing::TestWithParam<CodecType> {
 239  protected:
 240   void SetUp() override { codec_ = getCodec(GetParam()); }
 241
 242   void runSimpleTest(const DataHolder& dh);
 243
 244   std::unique_ptr<Codec> codec_;
 245 };
 246
 247 void CompressionCorruptionTest::runSimpleTest(const DataHolder& dh) {
 248   constexpr uint64_t uncompressedLength = 42;
 249   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength));
 250   auto compressed = codec_->compress(original.get());
 251
 252   if (!codec_->needsUncompressedLength()) {
 253     auto uncompressed = codec_->uncompress(compressed.get());
 254     EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
 255     EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
 256   }
 257   {
 258     auto uncompressed = codec_->uncompress(compressed.get(),
 259                                            uncompressedLength);
 260     EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
 261     EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
 262   }
 263
 264   EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength + 1),
 265                std::runtime_error);
 266
 267   // Corrupt the first character
 268   ++(compressed->writableData()[0]);
 269
 270   if (!codec_->needsUncompressedLength()) {
 271     EXPECT_THROW(codec_->uncompress(compressed.get()),
 272                  std::runtime_error);
 273   }
 274
 275   EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength),
 276                std::runtime_error);
 277 }
 278
 279 TEST_P(CompressionCorruptionTest, RandomData) {
 280   runSimpleTest(randomDataHolder);
 281 }
 282
 283 TEST_P(CompressionCorruptionTest, ConstantData) {
 284   runSimpleTest(constantDataHolder);
 285 }
 286
 287 INSTANTIATE_TEST_CASE_P(
 288     CompressionCorruptionTest,
 289     CompressionCorruptionTest,
 290     testing::Values(
 291         // NO_COMPRESSION can't detect corruption
 292         // LZ4 can't detect corruption reliably (sigh)
 293         CodecType::SNAPPY,
 294         CodecType::ZLIB));
 295
 296 }}}  // namespaces
 297
 298 int main(int argc, char *argv[]) {
 299   testing::InitGoogleTest(&argc, argv);
 300   gflags::ParseCommandLineFlags(&argc, &argv, true);
 301
 302   auto ret = RUN_ALL_TESTS();
 303   if (!ret) {
 304     folly::runBenchmarksOnFlag();
 305   }
 306   return ret;
 307 }