2 * Copyright 2013 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include "folly/io/Compression.h"
19 // Yes, tr1, as that's what gtest requires
23 #include <unordered_map>
25 #include <glog/logging.h>
26 #include <gtest/gtest.h>
28 #include "folly/Benchmark.h"
29 #include "folly/Hash.h"
30 #include "folly/Random.h"
31 #include "folly/io/IOBufQueue.h"
33 namespace folly { namespace io { namespace test {
35 constexpr size_t randomDataSizeLog2 = 27; // 128MiB
36 constexpr size_t randomDataSize = size_t(1) << randomDataSizeLog2;
38 std::unique_ptr<uint8_t[]> randomData;
39 std::unordered_map<uint64_t, uint64_t> hashes;
41 uint64_t hashIOBuf(const IOBuf* buf) {
42 uint64_t h = folly::hash::FNV_64_HASH_START;
43 for (auto& range : *buf) {
44 h = folly::hash::fnv64_buf(range.data(), range.size(), h);
49 uint64_t getRandomDataHash(uint64_t size) {
50 auto p = hashes.find(size);
51 if (p != hashes.end()) {
55 uint64_t h = folly::hash::fnv64_buf(randomData.get(), size);
60 void generateRandomData() {
61 randomData.reset(new uint8_t[size_t(1) << randomDataSizeLog2]);
63 constexpr size_t numThreadsLog2 = 3;
64 constexpr size_t numThreads = size_t(1) << numThreadsLog2;
66 uint32_t seed = randomNumberSeed();
68 std::vector<std::thread> threads;
69 threads.reserve(numThreads);
70 for (size_t t = 0; t < numThreads; ++t) {
72 [seed, t, numThreadsLog2] () {
73 std::mt19937 rng(seed + t);
74 size_t countLog2 = size_t(1) << (randomDataSizeLog2 - numThreadsLog2);
75 size_t start = size_t(t) << countLog2;
76 for (size_t i = 0; i < countLog2; ++i) {
77 randomData[start + i] = rng();
82 for (auto& t : threads) {
87 TEST(CompressionTestNeedsUncompressedLength, Simple) {
88 EXPECT_FALSE(getCodec(CodecType::NO_COMPRESSION)->needsUncompressedLength());
89 EXPECT_TRUE(getCodec(CodecType::LZ4)->needsUncompressedLength());
90 EXPECT_FALSE(getCodec(CodecType::SNAPPY)->needsUncompressedLength());
91 EXPECT_FALSE(getCodec(CodecType::ZLIB)->needsUncompressedLength());
92 EXPECT_FALSE(getCodec(CodecType::LZ4_VARINT_SIZE)->needsUncompressedLength());
95 class CompressionTest : public testing::TestWithParam<
96 std::tr1::tuple<int, CodecType>> {
99 auto tup = GetParam();
100 uncompressedLength_ = uint64_t(1) << std::tr1::get<0>(tup);
101 codec_ = getCodec(std::tr1::get<1>(tup));
104 uint64_t uncompressedLength_;
105 std::unique_ptr<Codec> codec_;
108 TEST_P(CompressionTest, Simple) {
109 auto original = IOBuf::wrapBuffer(randomData.get(), uncompressedLength_);
110 auto compressed = codec_->compress(original.get());
111 if (!codec_->needsUncompressedLength()) {
112 auto uncompressed = codec_->uncompress(compressed.get());
113 EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
114 EXPECT_EQ(getRandomDataHash(uncompressedLength_),
115 hashIOBuf(uncompressed.get()));
118 auto uncompressed = codec_->uncompress(compressed.get(),
119 uncompressedLength_);
120 EXPECT_EQ(uncompressedLength_, uncompressed->computeChainDataLength());
121 EXPECT_EQ(getRandomDataHash(uncompressedLength_),
122 hashIOBuf(uncompressed.get()));
126 INSTANTIATE_TEST_CASE_P(
130 testing::Values(0, 1, 12, 22, int(randomDataSizeLog2)),
131 testing::Values(CodecType::NO_COMPRESSION,
135 CodecType::LZ4_VARINT_SIZE)));
137 class CompressionCorruptionTest : public testing::TestWithParam<CodecType> {
140 codec_ = getCodec(GetParam());
143 std::unique_ptr<Codec> codec_;
146 TEST_P(CompressionCorruptionTest, Simple) {
147 constexpr uint64_t uncompressedLength = 42;
148 auto original = IOBuf::wrapBuffer(randomData.get(), uncompressedLength);
149 auto compressed = codec_->compress(original.get());
151 if (!codec_->needsUncompressedLength()) {
152 auto uncompressed = codec_->uncompress(compressed.get());
153 EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
154 EXPECT_EQ(getRandomDataHash(uncompressedLength),
155 hashIOBuf(uncompressed.get()));
158 auto uncompressed = codec_->uncompress(compressed.get(),
160 EXPECT_EQ(uncompressedLength, uncompressed->computeChainDataLength());
161 EXPECT_EQ(getRandomDataHash(uncompressedLength),
162 hashIOBuf(uncompressed.get()));
165 EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength + 1),
168 // Corrupt the first character
169 ++(compressed->writableData()[0]);
171 if (!codec_->needsUncompressedLength()) {
172 EXPECT_THROW(codec_->uncompress(compressed.get()),
176 EXPECT_THROW(codec_->uncompress(compressed.get(), uncompressedLength),
180 INSTANTIATE_TEST_CASE_P(
181 CompressionCorruptionTest,
182 CompressionCorruptionTest,
184 // NO_COMPRESSION can't detect corruption
185 // LZ4 can't detect corruption reliably (sigh)
191 int main(int argc, char *argv[]) {
192 testing::InitGoogleTest(&argc, argv);
193 google::ParseCommandLineFlags(&argc, &argv, true);
195 folly::io::test::generateRandomData(); // 4GB
197 auto ret = RUN_ALL_TESTS();
199 folly::runBenchmarksOnFlag();