Add LZ4_FRAME codec
[folly.git] / folly / io / Compression.cpp
index aaed2374b83d545497d695356a7ea5abf22d5794..e0beac2a6720e65d87afd1529c9fea5b80ae6cfa 100644 (file)
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 
 #if FOLLY_HAVE_LIBLZ4
 #include <lz4.h>
+#include <lz4frame.h>
 #include <lz4hc.h>
 #endif
 
@@ -56,13 +57,26 @@ std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
   uint64_t len = data->computeChainDataLength();
   if (len == 0) {
     return IOBuf::create(0);
-  } else if (len > maxUncompressedLength()) {
+  }
+  if (len > maxUncompressedLength()) {
     throw std::runtime_error("Codec: uncompressed length too large");
   }
 
   return doCompress(data);
 }
 
+std::string Codec::compress(const StringPiece data) {
+  const uint64_t len = data.size();
+  if (len == 0) {
+    return "";
+  }
+  if (len > maxUncompressedLength()) {
+    throw std::runtime_error("Codec: uncompressed length too large");
+  }
+
+  return doCompressString(data);
+}
+
 std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
                                          uint64_t uncompressedLength) {
   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
@@ -84,6 +98,28 @@ std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
   return doUncompress(data, uncompressedLength);
 }
 
+std::string Codec::uncompress(
+    const StringPiece data,
+    uint64_t uncompressedLength) {
+  if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
+    if (needsUncompressedLength()) {
+      throw std::invalid_argument("Codec: uncompressed length required");
+    }
+  } else if (uncompressedLength > maxUncompressedLength()) {
+    throw std::runtime_error("Codec: uncompressed length too large");
+  }
+
+  if (data.empty()) {
+    if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
+        uncompressedLength != 0) {
+      throw std::runtime_error("Codec: invalid uncompressed length");
+    }
+    return "";
+  }
+
+  return doUncompressString(data, uncompressedLength);
+}
+
 bool Codec::needsUncompressedLength() const {
   return doNeedsUncompressedLength();
 }
@@ -100,6 +136,30 @@ uint64_t Codec::doMaxUncompressedLength() const {
   return UNLIMITED_UNCOMPRESSED_LENGTH;
 }
 
+std::string Codec::doCompressString(const StringPiece data) {
+  const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
+  auto outputBuffer = doCompress(&inputBuffer);
+  std::string output;
+  output.reserve(outputBuffer->computeChainDataLength());
+  for (auto range : *outputBuffer) {
+    output.append(reinterpret_cast<const char*>(range.data()), range.size());
+  }
+  return output;
+}
+
+std::string Codec::doUncompressString(
+    const StringPiece data,
+    uint64_t uncompressedLength) {
+  const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
+  auto outputBuffer = doUncompress(&inputBuffer, uncompressedLength);
+  std::string output;
+  output.reserve(outputBuffer->computeChainDataLength());
+  for (auto range : *outputBuffer) {
+    output.append(reinterpret_cast<const char*>(range.data()), range.size());
+  }
+  return output;
+}
+
 namespace {
 
 /**
@@ -244,12 +304,11 @@ uint64_t LZ4Codec::doMaxUncompressedLength() const {
 }
 
 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
-  std::unique_ptr<IOBuf> clone;
+  IOBuf clone;
   if (data->isChained()) {
     // LZ4 doesn't support streaming, so we have to coalesce
-    clone = data->clone();
-    clone->coalesce();
-    data = clone.get();
+    clone = data->cloneCoalescedAsValue();
+    data = &clone;
   }
 
   uint32_t extraSize = encodeSize() ? kMaxVarintLength64 : 0;
@@ -286,12 +345,11 @@ std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
     const IOBuf* data,
     uint64_t uncompressedLength) {
-  std::unique_ptr<IOBuf> clone;
+  IOBuf clone;
   if (data->isChained()) {
     // LZ4 doesn't support streaming, so we have to coalesce
-    clone = data->clone();
-    clone->coalesce();
-    data = clone.get();
+    clone = data->cloneCoalescedAsValue();
+    data = &clone;
   }
 
   folly::io::Cursor cursor(data);
@@ -326,7 +384,160 @@ std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
   return out;
 }
 
-#endif  // FOLLY_HAVE_LIBLZ4
+class LZ4FrameCodec final : public Codec {
+ public:
+  static std::unique_ptr<Codec> create(int level, CodecType type);
+  explicit LZ4FrameCodec(int level, CodecType type);
+  ~LZ4FrameCodec();
+
+ private:
+  std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
+  std::unique_ptr<IOBuf> doUncompress(
+      const IOBuf* data,
+      uint64_t uncompressedLength) override;
+
+  // Reset the dctx_ if it is dirty or null.
+  void resetDCtx();
+
+  int level_;
+  LZ4F_dctx* dctx_{nullptr};
+  bool dirty_{false};
+};
+
+/* static */ std::unique_ptr<Codec> LZ4FrameCodec::create(
+    int level,
+    CodecType type) {
+  return make_unique<LZ4FrameCodec>(level, type);
+}
+
+static size_t lz4FrameThrowOnError(size_t code) {
+  if (LZ4F_isError(code)) {
+    throw std::runtime_error(
+        to<std::string>("LZ4Frame error: ", LZ4F_getErrorName(code)));
+  }
+  return code;
+}
+
+void LZ4FrameCodec::resetDCtx() {
+  if (dctx_ && !dirty_) {
+    return;
+  }
+  if (dctx_) {
+    LZ4F_freeDecompressionContext(dctx_);
+  }
+  lz4FrameThrowOnError(LZ4F_createDecompressionContext(&dctx_, 100));
+  dirty_ = false;
+}
+
+LZ4FrameCodec::LZ4FrameCodec(int level, CodecType type) : Codec(type) {
+  DCHECK(type == CodecType::LZ4_FRAME);
+  switch (level) {
+    case COMPRESSION_LEVEL_FASTEST:
+    case COMPRESSION_LEVEL_DEFAULT:
+      level_ = 0;
+      break;
+    case COMPRESSION_LEVEL_BEST:
+      level_ = 16;
+      break;
+    default:
+      level_ = level;
+      break;
+  }
+}
+
+LZ4FrameCodec::~LZ4FrameCodec() {
+  if (dctx_) {
+    LZ4F_freeDecompressionContext(dctx_);
+  }
+}
+
+std::unique_ptr<IOBuf> LZ4FrameCodec::doCompress(const IOBuf* data) {
+  // LZ4 Frame compression doesn't support streaming so we have to coalesce
+  IOBuf clone;
+  if (data->isChained()) {
+    clone = data->cloneCoalescedAsValue();
+    data = &clone;
+  }
+  // Set preferences
+  const auto uncompressedLength = data->length();
+  LZ4F_preferences_t prefs{};
+  prefs.compressionLevel = level_;
+  prefs.frameInfo.contentSize = uncompressedLength;
+  // Compress
+  auto buf = IOBuf::create(LZ4F_compressFrameBound(uncompressedLength, &prefs));
+  const size_t written = lz4FrameThrowOnError(LZ4F_compressFrame(
+      buf->writableTail(),
+      buf->tailroom(),
+      data->data(),
+      data->length(),
+      &prefs));
+  buf->append(written);
+  return buf;
+}
+
+std::unique_ptr<IOBuf> LZ4FrameCodec::doUncompress(
+    const IOBuf* data,
+    uint64_t uncompressedLength) {
+  // Reset the dctx if any errors have occurred
+  resetDCtx();
+  // Coalesce the data
+  ByteRange in = *data->begin();
+  IOBuf clone;
+  if (data->isChained()) {
+    clone = data->cloneCoalescedAsValue();
+    in = clone.coalesce();
+  }
+  data = nullptr;
+  // Select decompression options
+  LZ4F_decompressOptions_t options;
+  options.stableDst = 1;
+  // Select blockSize and growthSize for the IOBufQueue
+  IOBufQueue queue(IOBufQueue::cacheChainLength());
+  auto blockSize = uint64_t{64} << 10;
+  auto growthSize = uint64_t{4} << 20;
+  if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH) {
+    // Allocate uncompressedLength in one chunk (up to 64 MB)
+    const auto allocateSize = std::min(uncompressedLength, uint64_t{64} << 20);
+    queue.preallocate(allocateSize, allocateSize);
+    blockSize = std::min(uncompressedLength, blockSize);
+    growthSize = std::min(uncompressedLength, growthSize);
+  } else {
+    // Reduce growthSize for small data
+    const auto guessUncompressedLen = 4 * std::max(blockSize, in.size());
+    growthSize = std::min(guessUncompressedLen, growthSize);
+  }
+  // Once LZ4_decompress() is called, the dctx_ cannot be reused until it
+  // returns 0
+  dirty_ = true;
+  // Decompress until the frame is over
+  size_t code = 0;
+  do {
+    // Allocate enough space to decompress at least a block
+    void* out;
+    size_t outSize;
+    std::tie(out, outSize) = queue.preallocate(blockSize, growthSize);
+    // Decompress
+    size_t inSize = in.size();
+    code = lz4FrameThrowOnError(
+        LZ4F_decompress(dctx_, out, &outSize, in.data(), &inSize, &options));
+    if (in.empty() && outSize == 0 && code != 0) {
+      // We passed no input, no output was produced, and the frame isn't over
+      // No more forward progress is possible
+      throw std::runtime_error("LZ4Frame error: Incomplete frame");
+    }
+    in.uncheckedAdvance(inSize);
+    queue.postallocate(outSize);
+  } while (code != 0);
+  // At this point the decompression context can be reused
+  dirty_ = false;
+  if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
+      queue.chainLength() != uncompressedLength) {
+    throw std::runtime_error("LZ4Frame error: Invalid uncompressedLength");
+  }
+  return queue.move();
+}
+
+#endif // FOLLY_HAVE_LIBLZ4
 
 #if FOLLY_HAVE_LIBSNAPPY
 
@@ -1132,63 +1343,77 @@ std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
 
 }  // namespace
 
-std::unique_ptr<Codec> getCodec(CodecType type, int level) {
-  typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
-
-  static CodecFactory codecFactories[
-    static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
-    nullptr,  // USER_DEFINED
-    NoCompressionCodec::create,
+typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
+static constexpr CodecFactory
+    codecFactories[static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
+        nullptr, // USER_DEFINED
+        NoCompressionCodec::create,
 
 #if FOLLY_HAVE_LIBLZ4
-    LZ4Codec::create,
+        LZ4Codec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBSNAPPY
-    SnappyCodec::create,
+        SnappyCodec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBZ
-    ZlibCodec::create,
+        ZlibCodec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBLZ4
-    LZ4Codec::create,
+        LZ4Codec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBLZMA
-    LZMA2Codec::create,
-    LZMA2Codec::create,
+        LZMA2Codec::create,
+        LZMA2Codec::create,
 #else
-    nullptr,
-    nullptr,
+        nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBZSTD
-    ZSTDCodec::create,
+        ZSTDCodec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
 
 #if FOLLY_HAVE_LIBZ
-    ZlibCodec::create,
+        ZlibCodec::create,
 #else
-    nullptr,
+        nullptr,
 #endif
-  };
 
+#if FOLLY_HAVE_LIBLZ4
+        LZ4FrameCodec::create,
+#else
+        nullptr,
+#endif
+};
+
+bool hasCodec(CodecType type) {
   size_t idx = static_cast<size_t>(type);
   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
-    throw std::invalid_argument(to<std::string>(
-        "Compression type ", idx, " not supported"));
+    throw std::invalid_argument(
+        to<std::string>("Compression type ", idx, " invalid"));
+  }
+  return codecFactories[idx] != nullptr;
+}
+
+std::unique_ptr<Codec> getCodec(CodecType type, int level) {
+  size_t idx = static_cast<size_t>(type);
+  if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
+    throw std::invalid_argument(
+        to<std::string>("Compression type ", idx, " invalid"));
   }
   auto factory = codecFactories[idx];
   if (!factory) {