text: // Copyright (c) YugaByte, Inc.
  //
  // Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
  // in compliance with the License.  You may obtain a copy of the License at
  //
  // http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software distributed under the License
  // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  // or implied.  See the License for the specific language governing permissions and limitations
  // under the License.
  //
  
  #include "yb/rpc/compressed_stream.h"
  
  #include <lz4.h>
  #include <snappy-sinksource.h>
  #include <snappy.h>
  #include <zlib.h>
  
  #include <boost/preprocessor/cat.hpp>
  #include <boost/range/iterator_range.hpp>
  
  #include "yb/gutil/casts.h"
  
  #include "yb/rpc/circular_read_buffer.h"
  #include "yb/rpc/outbound_data.h"
  #include "yb/rpc/refined_stream.h"
  
  #include "yb/util/logging.h"
  #include "yb/util/result.h"
  #include "yb/util/size_literals.h"
  #include "yb/util/status_format.h"
  
  using namespace std::literals;
  
  DEFINE_int32(stream_compression_algo, 0, "Algorithm used for stream compression. "
                                           "0 - no compression, 1 - gzip, 2 - snappy, 3 - lz4.");
  
  namespace yb {
  namespace rpc {
  
  using SmallRefCntBuffers = boost::container::small_vector_base<RefCntBuffer>;
  
  namespace {
  
  class Compressor {
   public:
    virtual std::string ToString() const = 0;
  
    // Initialize compressor, required since we don't use exceptions to return error from ctor.
    virtual Status Init() = 0;
  
    // Compress specified vector of input buffers into single output buffer.
    virtual Status Compress(
        const SmallRefCntBuffers& input, RefinedStream* stream, OutboundDataPtr data) = 0;
  
    // Decompress specified input slice to specified output buffer.
    virtual Result<ReadBufferFull> Decompress(StreamReadBuffer* inp, StreamReadBuffer* out) = 0;
  
    // Connection header associated with this compressor.
    virtual OutboundDataPtr ConnectionHeader() = 0;
  
    virtual ~Compressor() = default;
  };
  
  size_t EntrySize(const RefCntBuffer& buffer) {
    return buffer.size();
  }
  
  size_t EntrySize(const iovec& iov) {
    return iov.iov_len;
  }
  
  char* EntryData(const RefCntBuffer& buffer) {
    return buffer.data();
  }
  
  char* EntryData(const iovec& iov) {
    return static_cast<char*>(iov.iov_base);
  }
  
  template <class Collection>
  size_t TotalLen(const Collection& input) {
    size_t result = 0;
    for (const auto& buf : input) {
      result += EntrySize(buf);
    }
    return result;
  }
  
  template <class Compressor>
  OutboundDataPtr GetConnectionHeader() {
    // Compressed stream header has signature YBx, where x - compressor identifier.
    static auto result = std::make_shared<StringOutboundData>(
        "YB"s + Compressor::kId, Compressor::kId + "ConnectionHeader"s);
    return result;
  }
  
  template <class SliceDecompressor>
  Result<ReadBufferFull> DecompressBySlices(
      StreamReadBuffer* inp, StreamReadBuffer* out, const SliceDecompressor& slice_decompressor) {
    size_t consumed = 0;
    auto out_vecs = VERIFY_RESULT(out->PrepareAppend());
    auto out_it = out_vecs.begin();
    size_t appended = 0;
  
    for (const auto& iov : inp->AppendedVecs()) {
      Slice slice(static_cast<char*>(iov.iov_base), iov.iov_len);
      for (;;) {
        if (out_it->iov_len == 0) {
          if (++out_it == out_vecs.end()) {
            break;
          }
        }
        size_t len = VERIFY_RESULT(slice_decompressor(&slice, out_it->iov_base, out_it->iov_len));
        appended += len;
        IoVecRemovePrefix(len, &*out_it);
        if (slice.empty()) {
          break;
        }
      }
      consumed += iov.iov_len - slice.size();
      if (!slice.empty()) {
        break;
      }
    }
    out->DataAppended(appended);
    inp->Consume(consumed, Slice());
    return ReadBufferFull(out->Full());
  }
  
  class ZlibCompressor : public Compressor {
   public:
    static const char kId = 'G';
    static const int kIndex = 1;
  
    explicit ZlibCompressor(MemTrackerPtr mem_tracker) {
    }
  
    ~ZlibCompressor() {
      if (deflate_inited_) {
        int res = deflateEnd(&deflate_stream_);
        LOG_IF(WARNING, res != Z_OK && res != Z_DATA_ERROR)
            << "Failed to destroy deflate stream: " << res;
      }
      if (inflate_inited_) {
        int res = inflateEnd(&inflate_stream_);
        LOG_IF(WARNING, res != Z_OK) << "Failed to destroy inflate stream: " << res;
      }
    }
  
    OutboundDataPtr ConnectionHeader() override {
      return GetConnectionHeader<ZlibCompressor>();
    }
  
    Status Init() override {
      memset(&deflate_stream_, 0, sizeof(deflate_stream_));
      int res = deflateInit(&deflate_stream_, /* level= */ Z_DEFAULT_COMPRESSION);
      if (res != Z_OK) {
        return STATUS_FORMAT(RuntimeError, "Cannot init deflate stream: $0", res);
      }
      deflate_inited_ = true;
  
      memset(&inflate_stream_, 0, sizeof(inflate_stream_));
      res = inflateInit(&inflate_stream_);
      if (res != Z_OK) {
        return STATUS_FORMAT(RuntimeError, "Cannot init inflate stream: $0", res);
      }
      inflate_inited_ = true;
  
      return Status::OK();
    }
  
    std::string ToString() const override {
      return "Zlib";
    }
  
    Status Compress(
        const SmallRefCntBuffers& input, RefinedStream* stream, OutboundDataPtr data) override {
      RefCntBuffer output(deflateBound(&deflate_stream_, TotalLen(input)));
      deflate_stream_.avail_out = static_cast<unsigned int>(output.size());
      deflate_stream_.next_out = output.udata();
  
      for (auto it = input.begin(); it != input.end();) {
        const auto& buf = *it++;
        deflate_stream_.next_in = const_cast<Bytef*>(buf.udata());
        deflate_stream_.avail_in = static_cast<unsigned int>(buf.size());
  
        for (;;) {
          auto res = deflate(&deflate_stream_, it == input.end() ? Z_PARTIAL_FLUSH : Z_NO_FLUSH);
          if (res == Z_STREAM_END) {
            if (deflate_stream_.avail_in != 0) {
              return STATUS_FORMAT(
                  RuntimeError, "Stream end when input data still available: $0",
                  deflate_stream_.avail_in);
            }
            break;
          }
          if (res != Z_OK) {
            return STATUS_FORMAT(RuntimeError, "Compression failed: $0", res);
          }
          if (deflate_stream_.avail_in == 0) {
            break;
          }
        }
      }
  
      output.Shrink(deflate_stream_.next_out - output.udata());
  
      // Send compressed data to underlying stream.
      return stream->SendToLower(std::make_shared<SingleBufferOutboundData>(
          std::move(output), std::move(data)));
    }
  
    Result<ReadBufferFull> Decompress(StreamReadBuffer* inp, StreamReadBuffer* out) override {
      return DecompressBySlices(
          inp, out, [this](Slice* input, void* out, size_t outlen) -> Result<size_t> {
        inflate_stream_.next_in = const_cast<Bytef*>(pointer_cast<const Bytef*>(input->data()));
        inflate_stream_.avail_in = narrow_cast<uInt>(input->size());
        inflate_stream_.next_out = static_cast<Bytef*>(out);
        inflate_stream_.avail_out = narrow_cast<uInt>(outlen);
  
        int res = inflate(&inflate_stream_, Z_NO_FLUSH);
        if (res != Z_OK && res != Z_BUF_ERROR) {
          return STATUS_FORMAT(RuntimeError, "Decompression failed: $0", res);
        }
  
        input->remove_prefix(input->size() - inflate_stream_.avail_in);
        return outlen - inflate_stream_.avail_out;
      });
    }
  
   private:
    z_stream deflate_stream_;
    z_stream inflate_stream_;
    bool deflate_inited_ = false;
    bool inflate_inited_ = false;
  };
  
  // Source implementation that provides input from range of buffers.
  template <class It>
  class RangeSource : public snappy::Source {
   public:
    explicit RangeSource(const It& begin, const It& end)
        : current_(begin), end_(end),
          available_(TotalLen(boost::make_iterator_range(begin, end))) {}
  
    void Limit(size_t value) {
      available_ = value;
    }
  
    size_t Available() const override {
      return available_;
    }
  
    const char* Peek(size_t* len) override {
      if (available_ == 0) {
        *len = 0;
        return nullptr;
      }
      *len = std::min(EntrySize(*current_) - current_pos_, available_);
      return EntryData(*current_) + current_pos_;
    }
  
    void Rewind(size_t n) {
      available_ += n;
      for (;;) {
        if (current_pos_ >= n) {
          current_pos_ -= n;
          break;
        }
        n -= current_pos_;
        --current_;
        current_pos_ = EntrySize(*current_);
      }
    }
  
    void Skip(size_t n) override {
      if (!n) {
        return;
      }
      available_ -= n;
      if ((current_pos_ += n) >= EntrySize(*current_)) {
        current_pos_ = 0;
        ++current_;
        DCHECK(available_ == 0 ? current_ == end_ : current_ < end_)
            << "Available: " << available_ << ", left buffers: " << std::distance(current_, end_)
            << ", n: " << n;
      }
    }
  
   private:
    It current_; // Current buffer that we are reading.
    It end_;
    size_t current_pos_ = 0; // Position in current buffer.
    size_t available_; // How many bytes left in all buffers.
  };
  
  // Sink implementation that provides output to io vecs.
  class IoVecsSink : public snappy::Sink {
   public:
    explicit IoVecsSink(IoVecs* out) : out_(out->begin()) {}
  
    size_t total_appended() const {
      return total_appended_;
    }
  
    void Append(const char* bytes, size_t n) override {
      total_appended_ += n;
  
      if (bytes == out_->iov_base) {
        if (out_->iov_len == n) {
          ++out_;
        } else {
          IoVecRemovePrefix(n, &*out_);
        }
        return;
      }
  
      while (n > 0) {
        size_t current_len = out_->iov_len;
        if (current_len >= n) {
          memcpy(out_->iov_base, bytes, n);
          IoVecRemovePrefix(n, &*out_);
          n = 0;
        } else {
          memcpy(out_->iov_base, bytes, current_len);
          bytes += current_len;
          n -= current_len;
          ++out_;
        }
      }
    }
  
    char* GetAppendBufferVariable(
        size_t min_size, size_t desired_size_hint, char* scratch,
        size_t scratch_size, size_t* allocated_size) override {
      if (min_size <= out_->iov_len) {
        *allocated_size = out_->iov_len;
        return static_cast<char*>(out_->iov_base);
      }
      return Sink::GetAppendBufferVariable(
          min_size, desired_size_hint, scratch, scratch_size, allocated_size);
    }
  
   private:
    IoVecs::iterator out_;
    size_t total_appended_ = 0;
  };
  
  // Binary search to find max value so that max_compressed_len(value) fits into header_len bytes.
  template <class F>
  static size_t FindMaxChunkSize(size_t header_len, const F& max_compressed_len) {
    size_t max_value = (1ULL << (8 * header_len)) - 1;
    size_t l = 1;
    size_t r = max_value;
    while (r > l) {
      size_t m = (l + r + 1) / 2;
      if (implicit_cast<size_t>(max_compressed_len(narrow_cast<int>(m))) > max_value) {
        r = m - 1;
      } else {
        l = m;
      }
    }
    return l;
  }
  
  // Snappy does not support stream compression.
  // So we have to add block len, to start decompression only from block start.
  // We use 2 bytes for block len, so have to limit uncompressed block so that in worst case
  // compressed block size would fit into 2 bytes.
  constexpr size_t kSnappyHeaderLen = 2;
  const size_t kSnappyMaxChunkSize = FindMaxChunkSize(kSnappyHeaderLen, &snappy::MaxCompressedLength);
  
  class SnappyCompressor : public Compressor {
   public:
    static constexpr char kId = 'S';
    static constexpr int kIndex = 2;
    static constexpr size_t kHeaderLen = kSnappyHeaderLen;
  
    explicit SnappyCompressor(MemTrackerPtr mem_tracker) {
    }
  
    ~SnappyCompressor() {
    }
  
    OutboundDataPtr ConnectionHeader() override {
      return GetConnectionHeader<SnappyCompressor>();
    }
  
    Status Init() override {
      return Status::OK();
    }
  
    std::string ToString() const override {
      return "Snappy";
    }
  
    Status Compress(
        const SmallRefCntBuffers& input, RefinedStream* stream, OutboundDataPtr data) override {
      RangeSource<SmallRefCntBuffers::const_iterator> source(input.begin(), input.end());
      auto input_size = source.Available();
      bool stop = false;
      while (!stop) {
        // Split input into chunks of size kSnappyMaxChunkSize or less.
        if (input_size > kSnappyMaxChunkSize) {
          source.Limit(kSnappyMaxChunkSize);
          input_size -= kSnappyMaxChunkSize;
        } else {
          source.Limit(input_size);
          stop = true;
        }
        RefCntBuffer output(kHeaderLen + snappy::MaxCompressedLength(source.Available()));
        snappy::UncheckedByteArraySink sink(output.data() + kHeaderLen);
        auto compressed_len = snappy::Compress(&source, &sink);
        BigEndian::Store16(output.data(), compressed_len);
        output.Shrink(kHeaderLen + compressed_len);
        RETURN_NOT_OK(stream->SendToLower(std::make_shared<SingleBufferOutboundData>(
            std::move(output),
            // We processed last buffer, attach data to it, so it will be notified when this buffer
            // is transferred.
            stop ? std::move(data) : nullptr)));
      }
      return Status::OK();
    }
  
    Result<ReadBufferFull> Decompress(StreamReadBuffer* inp, StreamReadBuffer* out) override {
      auto inp_vecs = inp->AppendedVecs();
      RangeSource<IoVecs::const_iterator> source(inp_vecs.begin(), inp_vecs.end());
  
      auto total_consumed = 0;
      auto out_vecs = VERIFY_RESULT(out->PrepareAppend());
      auto total_output = IoVecsFullSize(out_vecs);
      IoVecsSink sink(&out_vecs);
  
      size_t input_size = source.Available();
      auto result = ReadBufferFull::kFalse;
      while (total_consumed + kHeaderLen <= input_size) {
        source.Limit(input_size - total_consumed);
        // Fetch chunk len.
        size_t len = 0;
        auto data = source.Peek(&len);
        size_t size;
        if (len >= kHeaderLen) {
          // Size fully contained in the first buffer.
          size = BigEndian::Load16(data);
          source.Skip(kHeaderLen);
        } else {
          // Size distributed between 2 blocks.
          // Since blocks are not empty and size exactly 2, we should expect that first block
          // contains 1 byte of size.
          if (len != 1) {
            return STATUS(RuntimeError, "Expected to peek at least one byte");
          }
          char buf[kHeaderLen];
          buf[0] = *data;
          source.Skip(1);
          data = source.Peek(&len);
          buf[1] = *data;
          source.Skip(1);
          size = BigEndian::Load16(buf);
        }
        // Check whether we already received full chunk.
        VLOG_WITH_FUNC(4)
            << "Total consumed: " << total_consumed << ", size: " << size << ", input: "
            << input_size;
        if (total_consumed + kHeaderLen + size > input_size) {
          break;
        }
        source.Limit(size);
        uint32_t length = 0;
        auto old_available = source.Available();
        if (!snappy::GetUncompressedLength(&source, &length)) {
          return STATUS(RuntimeError, "GetUncompressedLength failed");
        }
  
        // Check if we have space for decompressed block.
        VLOG_WITH_FUNC(4)
            << "Total appended: " << sink.total_appended() << ", length: " << length
            << ", total_output: " << total_output;
        if (sink.total_appended() + length > total_output) {
          result = ReadBufferFull::kTrue;
          break;
        }
        // Rollback to state before GetUncompressedLength.
        source.Rewind(old_available - source.Available());
        if (!snappy::Uncompress(&source, &sink)) {
          return STATUS(RuntimeError, "Decompression failed");
        }
        total_consumed += kHeaderLen + size;
      }
  
      out->DataAppended(sink.total_appended());
      inp->Consume(total_consumed, Slice());
  
      return result;
    }
  };
  
  // LZ4 could compress/decompress only continuous block of memory into similar block.
  // So we do that same as for Snappy, but decompression is even more complex.
  constexpr size_t kLZ4HeaderLen = 2;
  const size_t kLZ4MaxChunkSize = FindMaxChunkSize(kLZ4HeaderLen, &LZ4_compressBound);
  const size_t kLZ4BufferSize = 64_KB;
  
  class LZ4DecompressState {
   public:
    LZ4DecompressState(char* input_buffer, char* output_buffer, Slice* prev_decompress_data_left)
        : input_buffer_(input_buffer), output_buffer_(output_buffer),
          prev_decompress_data_left_(prev_decompress_data_left) {}
  
    Result<ReadBufferFull> Execute(StreamReadBuffer* inp, StreamReadBuffer* out) {
      outvecs_ = VERIFY_RESULT(out->PrepareAppend());
      out_it_ = outvecs_.begin();
  
      // Check if we previously decompressed some data that did not fit into output buffer.
      // So copy it now. See DecompressChunk for details.
      while (!prev_decompress_data_left_->empty()) {
        auto len = std::min(prev_decompress_data_left_->size(), out_it_->iov_len);
        memcpy(out_it_->iov_base, prev_decompress_data_left_->data(), len);
        total_add_ += len;
        IoVecRemovePrefix(len, &*out_it_);
        prev_decompress_data_left_->remove_prefix(len);
        if (out_it_->iov_len == 0) {
          if (++out_it_ == outvecs_.end()) {
            out->DataAppended(total_add_);
            return ReadBufferFull(out->Full());
          }
        }
      }
  
      // Remaining piece of data in the previous vec.
      Slice prev_input_slice;
      for (const auto& input_vec : inp->AppendedVecs()) {
        Slice input_slice(static_cast<char*>(input_vec.iov_base), input_vec.iov_len);
        if (!prev_input_slice.empty()) {
          size_t chunk_size;
          if (prev_input_slice.size() >= kLZ4HeaderLen) {
            // Size fully contained in the previous block.
            chunk_size = BigEndian::Load16(prev_input_slice.data());
          } else if (prev_input_slice.size() + input_slice.size() < kLZ4HeaderLen) {
            // Did not receive header yet. So exit and wait for more data received.
            break;
          } else {
            // Size distributed between 2 blocks.
            char buf[kLZ4HeaderLen];
            memcpy(buf, prev_input_slice.data(), prev_input_slice.size());
            memcpy(buf + prev_input_slice.size(), input_slice.data(),
                   kLZ4HeaderLen - prev_input_slice.size());
            chunk_size = BigEndian::Load16(buf);
          }
          if (kLZ4HeaderLen + chunk_size > prev_input_slice.size() + input_slice.size()) {
            // Did not receive full block yet. So exit and wait for more data received.
            // TODO Here we rely on the fact that we use circular buffer and could have at most 2
            //      blocks. In general case chunk could be distributed across several buffers.
            break;
          }
          if (prev_input_slice.size() > kLZ4HeaderLen) {
            // Data is distributed between 2 buffers.
            // Have to copy it into separate input buffer so it will be a single continuous block.
            size_t size_in_prev_slice = prev_input_slice.size() - kLZ4HeaderLen;
            size_t size_in_current_slice = chunk_size - size_in_prev_slice;
            memcpy(input_buffer_, prev_input_slice.data() + kLZ4HeaderLen, size_in_prev_slice);
            memcpy(input_buffer_ + size_in_prev_slice, input_slice.data(), size_in_current_slice);
            RETURN_NOT_OK(DecompressChunk(Slice(input_buffer_, chunk_size)));
            input_slice.remove_prefix(size_in_current_slice);
          } else {
            // Only header (or part of header) was in previous buffer.
            // Could decompress from current buffer.
            input_slice.remove_prefix(kLZ4HeaderLen - prev_input_slice.size());
            RETURN_NOT_OK(DecompressChunk(input_slice.Prefix(chunk_size)));
            input_slice.remove_prefix(chunk_size);
          }
        }
        // Decompress all chunks contained in current buffer.
        while (input_slice.size() >= kLZ4HeaderLen && out_it_ != outvecs_.end()) {
          size_t chunk_size = BigEndian::Load16(input_slice.data());
          if (input_slice.size() < kLZ4HeaderLen + chunk_size) {
            break;
          }
          input_slice.remove_prefix(kLZ4HeaderLen);
          RETURN_NOT_OK(DecompressChunk(input_slice.Prefix(chunk_size)));
          input_slice.remove_prefix(chunk_size);
        }
        prev_input_slice = input_slice;
        // DecompressChunk could increase out_it_, so check whether we still have output space.
        if (out_it_ == outvecs_.end()) {
          break;
        }
      }
      inp->Consume(total_consumed_, Slice());
      out->DataAppended(total_add_);
      return ReadBufferFull(out->Full());
    }
  
   private:
    Status DecompressChunk(const Slice& input) {
      int res = LZ4_decompress_safe(
          input.cdata(), static_cast<char*>(out_it_->iov_base), narrow_cast<int>(input.size()),
          narrow_cast<int>(out_it_->iov_len));
      if (res <= 0) {
        // Unfortunately LZ4 does not provide information whether decryption failed because
        // of wrong data or it just does not fit into output buffer.
        // Try to decode to buffer that is big enough for max possible decompressed chunk.
        res = LZ4_decompress_safe(
            input.cdata(), output_buffer_, narrow_cast<int>(input.size()), kLZ4BufferSize);
        if (res <= 0) {
          return STATUS_FORMAT(RuntimeError, "Decompress failed: $0", res);
        }
  
        // Copy data from output buffer to provided read buffer.
        size_t size = res;
        char* buf = output_buffer_;
        while (out_it_ != outvecs_.end()) {
          size_t len = std::min<size_t>(size, out_it_->iov_len);
          memcpy(out_it_->iov_base, buf, len);
          IoVecRemovePrefix(len, &*out_it_);
          size -= len;
          total_add_ += len;
          buf += len;
          if (out_it_->iov_len == 0) {
            // Need to fill next output io vec.
            ++out_it_;
          }
          if (size == 0) {
            // Fully copied all decompressed data.
            break;
          }
        }
        if (size != 0) {
          // Have more decompressed data than provided read buffer could accept.
          *prev_decompress_data_left_ = Slice(buf, size);
        }
      } else {
        IoVecRemovePrefix(res, &*out_it_);
        total_add_ += res;
        if (out_it_->iov_len == 0) {
          ++out_it_;
        }
      }
  
      // Header was consumed by the caller, so we also add it here.
      total_consumed_ += kLZ4HeaderLen + input.size();
      return Status::OK();
    }
  
    // LZ4 operates on continuous chunks of memory, so we use input and output buffers to combine
    // several iovecs into one continuous chunk when necessary.
    char* input_buffer_;
    char* output_buffer_;
    Slice* prev_decompress_data_left_;
  
    IoVecs outvecs_;
    IoVecs::iterator out_it_;
    size_t total_add_ = 0;
    size_t total_consumed_ = 0;
  };
  
  class LZ4Compressor : public Compressor {
   public:
    static constexpr char kId = 'L';
    static constexpr int kIndex = 3;
    static constexpr size_t kHeaderLen = kLZ4HeaderLen;
  
    explicit LZ4Compressor(MemTrackerPtr mem_tracker) {
      if (mem_tracker) {
        consumption_ = ScopedTrackedConsumption(std::move(mem_tracker), 2 * kLZ4BufferSize);
      }
    }
  
    OutboundDataPtr ConnectionHeader() override {
      return GetConnectionHeader<LZ4Compressor>();
    }
  
    Status Init() override {
      return Status::OK();
    }
  
    std::string ToString() const override {
      return "LZ4";
    }
  
    Status Compress(
        const SmallRefCntBuffers& input, RefinedStream* stream, OutboundDataPtr data) override {
      // Increment iterator in loop body to be able to check whether it is last iteration or not.
      for (auto input_it = input.begin(); input_it != input.end();) {
        Slice input_slice = input_it->AsSlice();
        ++input_it;
        while (!input_slice.empty()) {
          Slice chunk;
          // Split input into chunks of size kLZ4MaxChunkSize or less.
          if (input_slice.size() > kLZ4MaxChunkSize) {
            chunk = input_slice.Prefix(kLZ4MaxChunkSize);
          } else {
            chunk = input_slice;
          }
          input_slice.remove_prefix(chunk.size());
          RefCntBuffer output(kHeaderLen + LZ4_compressBound(narrow_cast<int>(chunk.size())));
          int res = LZ4_compress(
              chunk.cdata(), output.data() + kHeaderLen, narrow_cast<int>(chunk.size()));
          if (res <= 0) {
            return STATUS_FORMAT(RuntimeError, "LZ4 compression failed: $0", res);
          }
          BigEndian::Store16(output.data(), res);
          output.Shrink(kHeaderLen + res);
          RETURN_NOT_OK(stream->SendToLower(std::make_shared<SingleBufferOutboundData>(
              std::move(output),
              // We processed last buffer, attach data to it, so it will be notified when this buffer
              // is transferred.
              input_slice.empty() && input_it == input.end() ? std::move(data) : nullptr)));
        }
      }
  
      return Status::OK();
    }
  
    Result<ReadBufferFull> Decompress(StreamReadBuffer* inp, StreamReadBuffer* out) override {
      LZ4DecompressState state(
          decompress_input_buf_, decompress_output_buf_, &prev_decompress_data_left_);
      return state.Execute(inp, out);
    }
  
   private:
    char decompress_input_buf_[kLZ4BufferSize];
    char decompress_output_buf_[kLZ4BufferSize];
    Slice prev_decompress_data_left_;
    ScopedTrackedConsumption consumption_;
  };
  
  #undef LZ4
  #define YB_COMPRESSION_ALGORITHMS (Zlib)(Snappy)(LZ4)
  
  #define YB_CREATE_COMPRESSOR_CASE(r, data, name) \
    case BOOST_PP_CAT(name, Compressor)::data: \
        return std::make_unique<BOOST_PP_CAT(name, Compressor)>(std::move(mem_tracker));
  
  std::unique_ptr<Compressor> CreateCompressor(char sign, MemTrackerPtr mem_tracker) {
    switch (sign) {
  BOOST_PP_SEQ_FOR_EACH(YB_CREATE_COMPRESSOR_CASE, kId, YB_COMPRESSION_ALGORITHMS)
      default:
        return nullptr;
    }
  }
  
  std::unique_ptr<Compressor> CreateOutboundCompressor(MemTrackerPtr mem_tracker) {
    auto algo = FLAGS_stream_compression_algo;
    if (!algo) {
      return nullptr;
    }
    switch (algo) {
  BOOST_PP_SEQ_FOR_EACH(YB_CREATE_COMPRESSOR_CASE, kIndex, YB_COMPRESSION_ALGORITHMS)
      default:
        YB_LOG_EVERY_N_SECS(DFATAL, 5) << "Unknown compression algorithm: " << algo;
        return nullptr;
    }
  }
  
  class CompressedRefiner : public StreamRefiner {
   public:
    CompressedRefiner() = default;
  
   private:
    void Start(RefinedStream* stream) override {
      stream_ = stream;
    }
  
    Status ProcessHeader() override {
      constexpr int kHeaderLen = 3;
  
      auto data = stream_->ReadBuffer().AppendedVecs();
      if (data.empty() || data[0].iov_len < kHeaderLen) {
        // Did not receive enough bytes to make a decision.
        // So just wait more bytes.
        return Status::OK();
      }
  
      const auto* bytes = static_cast<const uint8_t*>(data[0].iov_base);
      if (bytes[0] == 'Y' && bytes[1] == 'B') {
        compressor_ = CreateCompressor(bytes[2], stream_->buffer_tracker());
        if (compressor_) {
          RETURN_NOT_OK(compressor_->Init());
          RETURN_NOT_OK(stream_->StartHandshake());
          stream_->ReadBuffer().Consume(kHeaderLen, Slice());
          return Status::OK();
        }
      }
  
      // Don't use compression on this stream.
      return stream_->Established(RefinedStreamState::kDisabled);
    }
  
    Status Send(OutboundDataPtr data) override {
      boost::container::small_vector<RefCntBuffer, 10> input;
      data->Serialize(&input);
      return compressor_->Compress(input, stream_, std::move(data));
    }
  
    Status Handshake() override {
      if (stream_->local_side() == LocalSide::kClient) {
        compressor_ = CreateOutboundCompressor(stream_->buffer_tracker());
        if (!compressor_) {
          return stream_->Established(RefinedStreamState::kDisabled);
        }
        RETURN_NOT_OK(compressor_->Init());
        RETURN_NOT_OK(stream_->SendToLower(compressor_->ConnectionHeader()));
      }
  
      return stream_->Established(RefinedStreamState::kEnabled);
    }
  
    Result<ReadBufferFull> Read(StreamReadBuffer* out) override {
      VLOG_WITH_PREFIX(4) << __func__;
  
      return compressor_->Decompress(&stream_->ReadBuffer(), out);
    }
  
    const Protocol* GetProtocol() override {
      return CompressedStreamProtocol();
    }
  
    std::string ToString() const override {
      return compressor_ ? compressor_->ToString() : "PLAIN";
    }
  
    const std::string& LogPrefix() const {
      return stream_->LogPrefix();
    }
  
    RefinedStream* stream_ = nullptr;
    std::unique_ptr<Compressor> compressor_ = nullptr;
  };
  
  } // namespace
  
  const Protocol* CompressedStreamProtocol() {
    static Protocol result("tcpc");
    return &result;
  }
  
  StreamFactoryPtr CompressedStreamFactory(
      StreamFactoryPtr lower_layer_factory, const MemTrackerPtr& buffer_tracker) {
    return std::make_shared<RefinedStreamFactory>(
        std::move(lower_layer_factory), buffer_tracker, [](const StreamCreateData& data) {
      return std::make_unique<CompressedRefiner>();
    });
  }
  
  }  // namespace rpc
  }  // namespace yb
  