Commit 7dd7df55 authored by Nick Terrell's avatar Nick Terrell Committed by Facebook GitHub Bot

Surrender the ZSTD_{C,D}Ctx on (de)compression end

Summary: We already cache the contexts in the context pool. We don't also need to cache the contexts in the `Codec` when the user keeps the codec around. This will allow us to better utilize the contexts, and potentially reduce memory usage.

Reviewed By: felixhandte

Differential Revision: D26356275

fbshipit-source-id: 6d80326f469de61094478c73b6fe557eebcd7a22
parent bdd1109c
...@@ -126,7 +126,6 @@ class ZSTDStreamCodec final : public StreamCodec { ...@@ -126,7 +126,6 @@ class ZSTDStreamCodec final : public StreamCodec {
void resetDCtx(); void resetDCtx();
Options options_; Options options_;
bool needReset_{true};
ZSTD_CCtx_Pool::Ref cctx_{getNULL_ZSTD_CCtx()}; ZSTD_CCtx_Pool::Ref cctx_{getNULL_ZSTD_CCtx()};
ZSTD_DCtx_Pool::Ref dctx_{getNULL_ZSTD_DCtx()}; ZSTD_DCtx_Pool::Ref dctx_{getNULL_ZSTD_DCtx()};
}; };
...@@ -177,16 +176,14 @@ Optional<uint64_t> ZSTDStreamCodec::doGetUncompressedLength( ...@@ -177,16 +176,14 @@ Optional<uint64_t> ZSTDStreamCodec::doGetUncompressedLength(
} }
void ZSTDStreamCodec::doResetStream() { void ZSTDStreamCodec::doResetStream() {
needReset_ = true; cctx_.reset(nullptr);
dctx_.reset(nullptr);
} }
void ZSTDStreamCodec::resetCCtx() { void ZSTDStreamCodec::resetCCtx() {
if (!cctx_) { DCHECK(cctx_ == nullptr);
cctx_ = getZSTD_CCtx(); cctx_ = getZSTD_CCtx();
if (!cctx_) { DCHECK(cctx_ != nullptr);
throw std::bad_alloc{};
}
}
resetCCtxSessionAndParameters(cctx_.get()); resetCCtxSessionAndParameters(cctx_.get());
zstdThrowIfError( zstdThrowIfError(
ZSTD_CCtx_setParametersUsingCCtxParams(cctx_.get(), options_.params())); ZSTD_CCtx_setParametersUsingCCtxParams(cctx_.get(), options_.params()));
...@@ -196,9 +193,8 @@ void ZSTDStreamCodec::resetCCtx() { ...@@ -196,9 +193,8 @@ void ZSTDStreamCodec::resetCCtx() {
bool ZSTDStreamCodec::doCompressStream( bool ZSTDStreamCodec::doCompressStream(
ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp flushOp) { ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp flushOp) {
if (needReset_) { if (cctx_ == nullptr) {
resetCCtx(); resetCCtx();
needReset_ = false;
} }
ZSTD_inBuffer in = {input.data(), input.size(), 0}; ZSTD_inBuffer in = {input.data(), input.size(), 0};
ZSTD_outBuffer out = {output.data(), output.size(), 0}; ZSTD_outBuffer out = {output.data(), output.size(), 0};
...@@ -212,7 +208,12 @@ bool ZSTDStreamCodec::doCompressStream( ...@@ -212,7 +208,12 @@ bool ZSTDStreamCodec::doCompressStream(
case StreamCodec::FlushOp::NONE: case StreamCodec::FlushOp::NONE:
return false; return false;
case StreamCodec::FlushOp::FLUSH: case StreamCodec::FlushOp::FLUSH:
return rc == 0;
case StreamCodec::FlushOp::END: case StreamCodec::FlushOp::END:
if (rc == 0) {
// Surrender our cctx_
doResetStream();
}
return rc == 0; return rc == 0;
default: default:
throw std::invalid_argument("ZSTD: invalid FlushOp"); throw std::invalid_argument("ZSTD: invalid FlushOp");
...@@ -220,12 +221,9 @@ bool ZSTDStreamCodec::doCompressStream( ...@@ -220,12 +221,9 @@ bool ZSTDStreamCodec::doCompressStream(
} }
void ZSTDStreamCodec::resetDCtx() { void ZSTDStreamCodec::resetDCtx() {
if (!dctx_) { DCHECK(dctx_ == nullptr);
dctx_ = getZSTD_DCtx(); dctx_ = getZSTD_DCtx();
if (!dctx_) { DCHECK(dctx_ != nullptr);
throw std::bad_alloc{};
}
}
resetDCtxSessionAndParameters(dctx_.get()); resetDCtxSessionAndParameters(dctx_.get());
if (options_.maxWindowSize() != 0) { if (options_.maxWindowSize() != 0) {
zstdThrowIfError( zstdThrowIfError(
...@@ -235,9 +233,8 @@ void ZSTDStreamCodec::resetDCtx() { ...@@ -235,9 +233,8 @@ void ZSTDStreamCodec::resetDCtx() {
bool ZSTDStreamCodec::doUncompressStream( bool ZSTDStreamCodec::doUncompressStream(
ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp) { ByteRange& input, MutableByteRange& output, StreamCodec::FlushOp) {
if (needReset_) { if (dctx_ == nullptr) {
resetDCtx(); resetDCtx();
needReset_ = false;
} }
ZSTD_inBuffer in = {input.data(), input.size(), 0}; ZSTD_inBuffer in = {input.data(), input.size(), 0};
ZSTD_outBuffer out = {output.data(), output.size(), 0}; ZSTD_outBuffer out = {output.data(), output.size(), 0};
...@@ -247,6 +244,10 @@ bool ZSTDStreamCodec::doUncompressStream( ...@@ -247,6 +244,10 @@ bool ZSTDStreamCodec::doUncompressStream(
}; };
size_t const rc = size_t const rc =
zstdThrowIfError(ZSTD_decompressStream(dctx_.get(), &out, &in)); zstdThrowIfError(ZSTD_decompressStream(dctx_.get(), &out, &in));
if (rc == 0) {
// Surrender our dctx_
doResetStream();
}
return rc == 0; return rc == 0;
} }
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment