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   EXPECT_TRUE(getCodec(CodecType::ZSTD_BETA)->needsUncompressedLength());
 132 }
 133
 134 class CompressionTest
 135     : public testing::TestWithParam<std::tr1::tuple<int, CodecType>> {
 136   protected:
 137    void SetUp() override {
 138      auto tup = GetParam();
 139      uncompressedLength_ = uint64_t(1) << std::tr1::get<0>(tup);
 140      codec_ = getCodec(std::tr1::get<1>(tup));
 141    }
 142
 143    void runSimpleTest(const DataHolder& dh);
 144
 145    uint64_t uncompressedLength_;
 146    std::unique_ptr<Codec> codec_;
 147 };
 148
 149 void CompressionTest::runSimpleTest(const DataHolder& dh) {
 150   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength_));
 151   auto compressed = codec_->compress(original.get());
 152   if (!codec_->needsUncompressedLength()) {
 153     auto uncompressed = codec_->uncompress(compressed.get());
 154
 155     EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 156     EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 157   }
 158   {
 159     auto uncompressed = codec_->uncompress(compressed.get(),
 160                                            uncompressedLength_);
 161     EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 162     EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 163   }
 164 }
 165
 166 TEST_P(CompressionTest, RandomData) {
 167   runSimpleTest(randomDataHolder);
 168 }
 169
 170 TEST_P(CompressionTest, ConstantData) {
 171   runSimpleTest(constantDataHolder);
 172 }
 173
 174 INSTANTIATE_TEST_CASE_P(
 175     CompressionTest,
 176     CompressionTest,
 177     testing::Combine(testing::Values(0, 1, 12, 22, 25, 27),
 178                      testing::Values(CodecType::NO_COMPRESSION,
 179                                      CodecType::LZ4,
 180                                      CodecType::SNAPPY,
 181                                      CodecType::ZLIB,
 182                                      CodecType::LZ4_VARINT_SIZE,
 183                                      CodecType::LZMA2,
 184                                      CodecType::LZMA2_VARINT_SIZE,
 185                                      CodecType::ZSTD_BETA)));
 186
 187 class CompressionVarintTest
 188     : public testing::TestWithParam<std::tr1::tuple<int, CodecType>> {
 189  protected:
 190   void SetUp() override {
 191     auto tup = GetParam();
 192     uncompressedLength_ = uint64_t(1) << std::tr1::get<0>(tup);
 193     codec_ = getCodec(std::tr1::get<1>(tup));
 194   }
 195
 196   void runSimpleTest(const DataHolder& dh);
 197
 198   uint64_t uncompressedLength_;
 199   std::unique_ptr<Codec> codec_;
 200 };
 201
 202 inline uint64_t oneBasedMsbPos(uint64_t number) {
 203   uint64_t pos = 0;
 204   for (; number > 0; ++pos, number >>= 1) {
 205   }
 206   return pos;
 207 }
 208
 209 void CompressionVarintTest::runSimpleTest(const DataHolder& dh) {
 210   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength_));
 211   auto compressed = codec_->compress(original.get());
 212   auto breakPoint =
 213       1UL +
 214       Random::rand64(std::max(9UL, oneBasedMsbPos(uncompressedLength_)) / 9UL);
 215   auto tinyBuf = IOBuf::copyBuffer(compressed->data(),
 216                                    std::min(compressed->length(), breakPoint));
 217   compressed->trimStart(breakPoint);
 218   tinyBuf->prependChain(std::move(compressed));
 219   compressed = std::move(tinyBuf);
 220
 221   auto uncompressed = codec_->uncompress(compressed.get());
 222
 223   EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
 224   EXPECT_EQ(dh.hash(uncompressedLength_), hashIOBuf(uncompressed.get()));
 225 }
 226
 227 TEST_P(CompressionVarintTest, RandomData) { runSimpleTest(randomDataHolder); }
 228
 229 TEST_P(CompressionVarintTest, ConstantData) {
 230   runSimpleTest(constantDataHolder);
 231 }
 232
 233 INSTANTIATE_TEST_CASE_P(
 234     CompressionVarintTest,
 235     CompressionVarintTest,
 236     testing::Combine(testing::Values(0, 1, 12, 22, 25, 27),
 237                      testing::Values(CodecType::LZ4_VARINT_SIZE,
 238                                      CodecType::LZMA2_VARINT_SIZE)));
 239
 240 class CompressionCorruptionTest : public testing::TestWithParam<CodecType> {
 241  protected:
 242   void SetUp() override { codec_ = getCodec(GetParam()); }
 243
 244   void runSimpleTest(const DataHolder& dh);
 245
 246   std::unique_ptr<Codec> codec_;
 247 };
 248
 249 void CompressionCorruptionTest::runSimpleTest(const DataHolder& dh) {
 250   constexpr uint64_t uncompressedLength = 42;
 251   auto original = IOBuf::wrapBuffer(dh.data(uncompressedLength));
 252   auto compressed = codec_->compress(original.get());
 253
 254   if (!codec_->needsUncompressedLength()) {
 255     auto uncompressed = codec_->uncompress(compressed.get());
 256     EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
 257     EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
 258   }
 259   {
 260     auto uncompressed = codec_->uncompress(compressed.get(),
 261                                            uncompressedLength);
 262     EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
 263     EXPECT_EQ(dh.hash(uncompressedLength), hashIOBuf(uncompressed.get()));
 264   }
 265
 266   EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength + 1),
 267                std::runtime_error);
 268
 269   // Corrupt the first character
 270   ++(compressed->writableData()[0]);
 271
 272   if (!codec_->needsUncompressedLength()) {
 273     EXPECT_THROW(codec_->uncompress(compressed.get()),
 274                  std::runtime_error);
 275   }
 276
 277   EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength),
 278                std::runtime_error);
 279 }
 280
 281 TEST_P(CompressionCorruptionTest, RandomData) {
 282   runSimpleTest(randomDataHolder);
 283 }
 284
 285 TEST_P(CompressionCorruptionTest, ConstantData) {
 286   runSimpleTest(constantDataHolder);
 287 }
 288
 289 INSTANTIATE_TEST_CASE_P(
 290     CompressionCorruptionTest,
 291     CompressionCorruptionTest,
 292     testing::Values(
 293         // NO_COMPRESSION can't detect corruption
 294         // LZ4 can't detect corruption reliably (sigh)
 295         CodecType::SNAPPY,
 296         CodecType::ZLIB));
 297
 298 }}}  // namespaces
 299
 300 int main(int argc, char *argv[]) {
 301   testing::InitGoogleTest(&argc, argv);
 302   gflags::ParseCommandLineFlags(&argc, &argv, true);
 303
 304   auto ret = RUN_ALL_TESTS();
 305   if (!ret) {
 306     folly::runBenchmarksOnFlag();
 307   }
 308   return ret;
 309 }