From fd893dfcb760fa097133c7c9bd964bd3e292fdbd Mon Sep 17 00:00:00 2001 From: Dean Lee Date: Tue, 17 Sep 2024 00:45:15 +0800 Subject: [PATCH] replay: optimize decompressZST with pre-allocated buffer (#33562) improve decompressZST --- tools/replay/util.cc | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tools/replay/util.cc b/tools/replay/util.cc index 81e87b7dd3..0ae0c05f9b 100644 --- a/tools/replay/util.cc +++ b/tools/replay/util.cc @@ -329,8 +329,17 @@ std::string decompressZST(const std::byte *in, size_t in_size, std::atomic // Initialize input and output buffers ZSTD_inBuffer input = {in, in_size, 0}; + + // Estimate and reserve memory for decompressed data + size_t estimatedDecompressedSize = ZSTD_getFrameContentSize(in, in_size); + if (estimatedDecompressedSize == ZSTD_CONTENTSIZE_ERROR || estimatedDecompressedSize == ZSTD_CONTENTSIZE_UNKNOWN) { + estimatedDecompressedSize = in_size * 2; // Use a fallback size + } + std::string decompressedData; - const size_t bufferSize = ZSTD_DStreamOutSize(); // recommended output buffer size + decompressedData.reserve(estimatedDecompressedSize); + + const size_t bufferSize = ZSTD_DStreamOutSize(); // Recommended output buffer size std::string outputBuffer(bufferSize, '\0'); while (input.pos < input.size && !(abort && *abort)) {