From 45b6bcca9877ffffe3c4d7386089f41a128d7256 Mon Sep 17 00:00:00 2001
From: Zhongyi Xie <xiez@fb.com>
Date: Wed, 6 Jun 2018 23:30:26 -0700
Subject: [PATCH] ZSTD compression: should also expect type =
 kZSTDNotFinalCompression (#3964)

Summary:
Depending on the compression type, `CompressBlock` calls the compress method for each compression type. It calls ZSTD_Compress for both kZSTD and kZSTDNotFinalCompression (https://github.com/facebook/rocksdb/blob/master/table/block_based_table_builder.cc#L169).
However currently ZSTD_Compress only expects the type to be kZSTD and this is causing assert failures and crashes. The same also applies to ZSTD_Uncompress.
Closes https://github.com/facebook/rocksdb/pull/3964

Differential Revision: D8308715

Pulled By: miasantreble

fbshipit-source-id: e5125f53edb829c9c33733167bec74e4793d0782
---
 util/compression.h | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/util/compression.h b/util/compression.h
index d6336b33f..af3fcf00e 100644
--- a/util/compression.h
+++ b/util/compression.h
@@ -140,7 +140,7 @@ class CompressionContext {
 #if defined(ZSTD) && (ZSTD_VERSION_NUMBER >= 500)
   ZSTD_CCtx* zstd_ctx_ = nullptr;
   void CreateNativeContext() {
-    if (type_ == kZSTD) {
+    if (type_ == kZSTD || type_ == kZSTDNotFinalCompression) {
 #ifdef ROCKSDB_ZSTD_CUSTOM_MEM
       zstd_ctx_ =
           ZSTD_createCCtx_advanced(port::GetJeZstdAllocationOverrides());
@@ -158,7 +158,7 @@ class CompressionContext {
  public:
   // callable inside ZSTD_Compress
   ZSTD_CCtx* ZSTDPreallocCtx() const {
-    assert(type_ == kZSTD);
+    assert(type_ == kZSTD || type_ == kZSTDNotFinalCompression);
     return zstd_ctx_;
   }
 #else   // ZSTD && (ZSTD_VERSION_NUMBER >= 500)
@@ -201,13 +201,14 @@ class UncompressionContext {
       : UncompressionContext(comp_type, Slice()) {}
   UncompressionContext(CompressionType comp_type, const Slice& comp_dict)
       : type_(comp_type), dict_(comp_dict) {
-    if (type_ == kZSTD) {
+    if (type_ == kZSTD || type_ == kZSTDNotFinalCompression) {
       ctx_cache_ = CompressionContextCache::Instance();
       uncomp_cached_data_ = ctx_cache_->GetCachedZSTDUncompressData();
     }
   }
   ~UncompressionContext() {
-    if (type_ == kZSTD && uncomp_cached_data_.GetCacheIndex() != -1) {
+    if ((type_ == kZSTD || type_ == kZSTDNotFinalCompression) &&
+        uncomp_cached_data_.GetCacheIndex() != -1) {
       assert(ctx_cache_ != nullptr);
       ctx_cache_->ReturnCachedZSTDUncompressData(
           uncomp_cached_data_.GetCacheIndex());