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
This commit is contained in:
Zhongyi Xie 2018-06-06 23:30:26 -07:00 committed by Facebook Github Bot
parent b73652169e
commit 45b6bcca98

View File

@ -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());