619 lines
21 KiB
C++
619 lines
21 KiB
C++
|
/*
|
||
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
||
|
|
* All rights reserved.
|
||
|
|
*
|
||
|
|
* This source code is licensed under both the BSD-style license (found in the
|
||
|
|
* LICENSE file in the root directory of this source tree) and the GPLv2 (found
|
||
|
|
* in the COPYING file in the root directory of this source tree).
|
||
|
|
*/
|
||
|
|
#include "platform.h" /* Large Files support, SET_BINARY_MODE */
|
||
|
|
#include "Pzstd.h"
|
||
|
|
#include "SkippableFrame.h"
|
||
|
|
#include "utils/FileSystem.h"
|
||
|
|
#include "utils/Portability.h"
|
||
|
|
#include "utils/Range.h"
|
||
|
|
#include "utils/ScopeGuard.h"
|
||
|
|
#include "utils/ThreadPool.h"
|
||
|
|
#include "utils/WorkQueue.h"
|
||
|
|
|
||
|
|
#include <algorithm>
|
||
|
|
#include <chrono>
|
||
|
|
#include <cinttypes>
|
||
|
|
#include <cstddef>
|
||
|
|
#include <cstdio>
|
||
|
|
#include <memory>
|
||
|
|
#include <string>
|
||
|
|
|
||
|
|
|
||
|
|
namespace pzstd {
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
#ifdef _WIN32
|
||
|
|
const std::string nullOutput = "nul";
|
||
|
|
#else
|
||
|
|
const std::string nullOutput = "/dev/null";
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
|
||
|
|
using std::size_t;
|
||
|
|
|
||
|
|
static std::uintmax_t fileSizeOrZero(const std::string &file) {
|
||
|
|
if (file == "-") {
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
std::error_code ec;
|
||
|
|
auto size = file_size(file, ec);
|
||
|
|
if (ec) {
|
||
|
|
size = 0;
|
||
|
|
}
|
||
|
|
return size;
|
||
|
|
}
|
||
|
|
|
||
|
|
static std::uint64_t handleOneInput(const Options &options,
|
||
|
|
const std::string &inputFile,
|
||
|
|
FILE* inputFd,
|
||
|
|
const std::string &outputFile,
|
||
|
|
FILE* outputFd,
|
||
|
|
SharedState& state) {
|
||
|
|
auto inputSize = fileSizeOrZero(inputFile);
|
||
|
|
// WorkQueue outlives ThreadPool so in the case of error we are certain
|
||
|
|
// we don't accidentally try to call push() on it after it is destroyed
|
||
|
|
WorkQueue<std::shared_ptr<BufferWorkQueue>> outs{options.numThreads + 1};
|
||
|
|
std::uint64_t bytesRead;
|
||
|
|
std::uint64_t bytesWritten;
|
||
|
|
{
|
||
|
|
// Initialize the (de)compression thread pool with numThreads
|
||
|
|
ThreadPool executor(options.numThreads);
|
||
|
|
// Run the reader thread on an extra thread
|
||
|
|
ThreadPool readExecutor(1);
|
||
|
|
if (!options.decompress) {
|
||
|
|
// Add a job that reads the input and starts all the compression jobs
|
||
|
|
readExecutor.add(
|
||
|
|
[&state, &outs, &executor, inputFd, inputSize, &options, &bytesRead] {
|
||
|
|
bytesRead = asyncCompressChunks(
|
||
|
|
state,
|
||
|
|
outs,
|
||
|
|
executor,
|
||
|
|
inputFd,
|
||
|
|
inputSize,
|
||
|
|
options.numThreads,
|
||
|
|
options.determineParameters());
|
||
|
|
});
|
||
|
|
// Start writing
|
||
|
|
bytesWritten = writeFile(state, outs, outputFd, options.decompress);
|
||
|
|
} else {
|
||
|
|
// Add a job that reads the input and starts all the decompression jobs
|
||
|
|
readExecutor.add([&state, &outs, &executor, inputFd, &bytesRead] {
|
||
|
|
bytesRead = asyncDecompressFrames(state, outs, executor, inputFd);
|
||
|
|
});
|
||
|
|
// Start writing
|
||
|
|
bytesWritten = writeFile(state, outs, outputFd, options.decompress);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (!state.errorHolder.hasError()) {
|
||
|
|
std::string inputFileName = inputFile == "-" ? "stdin" : inputFile;
|
||
|
|
std::string outputFileName = outputFile == "-" ? "stdout" : outputFile;
|
||
|
|
if (!options.decompress) {
|
||
|
|
double ratio = static_cast<double>(bytesWritten) /
|
||
|
|
static_cast<double>(bytesRead + !bytesRead);
|
||
|
|
state.log(kLogInfo, "%-20s :%6.2f%% (%6" PRIu64 " => %6" PRIu64
|
||
|
|
" bytes, %s)\n",
|
||
|
|
inputFileName.c_str(), ratio * 100, bytesRead, bytesWritten,
|
||
|
|
outputFileName.c_str());
|
||
|
|
} else {
|
||
|
|
state.log(kLogInfo, "%-20s: %" PRIu64 " bytes \n",
|
||
|
|
inputFileName.c_str(),bytesWritten);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return bytesWritten;
|
||
|
|
}
|
||
|
|
|
||
|
|
static FILE *openInputFile(const std::string &inputFile,
|
||
|
|
ErrorHolder &errorHolder) {
|
||
|
|
if (inputFile == "-") {
|
||
|
|
SET_BINARY_MODE(stdin);
|
||
|
|
return stdin;
|
||
|
|
}
|
||
|
|
// Check if input file is a directory
|
||
|
|
{
|
||
|
|
std::error_code ec;
|
||
|
|
if (is_directory(inputFile, ec)) {
|
||
|
|
errorHolder.setError("Output file is a directory -- ignored");
|
||
|
|
return nullptr;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
auto inputFd = std::fopen(inputFile.c_str(), "rb");
|
||
|
|
if (!errorHolder.check(inputFd != nullptr, "Failed to open input file")) {
|
||
|
|
return nullptr;
|
||
|
|
}
|
||
|
|
return inputFd;
|
||
|
|
}
|
||
|
|
|
||
|
|
static FILE *openOutputFile(const Options &options,
|
||
|
|
const std::string &outputFile,
|
||
|
|
SharedState& state) {
|
||
|
|
if (outputFile == "-") {
|
||
|
|
SET_BINARY_MODE(stdout);
|
||
|
|
return stdout;
|
||
|
|
}
|
||
|
|
// Check if the output file exists and then open it
|
||
|
|
if (!options.overwrite && outputFile != nullOutput) {
|
||
|
|
auto outputFd = std::fopen(outputFile.c_str(), "rb");
|
||
|
|
if (outputFd != nullptr) {
|
||
|
|
std::fclose(outputFd);
|
||
|
|
if (!state.log.logsAt(kLogInfo)) {
|
||
|
|
state.errorHolder.setError("Output file exists");
|
||
|
|
return nullptr;
|
||
|
|
}
|
||
|
|
state.log(
|
||
|
|
kLogInfo,
|
||
|
|
"pzstd: %s already exists; do you wish to overwrite (y/n) ? ",
|
||
|
|
outputFile.c_str());
|
||
|
|
int c = getchar();
|
||
|
|
if (c != 'y' && c != 'Y') {
|
||
|
|
state.errorHolder.setError("Not overwritten");
|
||
|
|
return nullptr;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
auto outputFd = std::fopen(outputFile.c_str(), "wb");
|
||
|
|
if (!state.errorHolder.check(
|
||
|
|
outputFd != nullptr, "Failed to open output file")) {
|
||
|
|
return nullptr;
|
||
|
|
}
|
||
|
|
return outputFd;
|
||
|
|
}
|
||
|
|
|
||
|
|
int pzstdMain(const Options &options) {
|
||
|
|
int returnCode = 0;
|
||
|
|
SharedState state(options);
|
||
|
|
for (const auto& input : options.inputFiles) {
|
||
|
|
// Setup the shared state
|
||
|
|
auto printErrorGuard = makeScopeGuard([&] {
|
||
|
|
if (state.errorHolder.hasError()) {
|
||
|
|
returnCode = 1;
|
||
|
|
state.log(kLogError, "pzstd: %s: %s.\n", input.c_str(),
|
||
|
|
state.errorHolder.getError().c_str());
|
||
|
|
}
|
||
|
|
});
|
||
|
|
// Open the input file
|
||
|
|
auto inputFd = openInputFile(input, state.errorHolder);
|
||
|
|
if (inputFd == nullptr) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
auto closeInputGuard = makeScopeGuard([&] { std::fclose(inputFd); });
|
||
|
|
// Open the output file
|
||
|
|
auto outputFile = options.getOutputFile(input);
|
||
|
|
if (!state.errorHolder.check(outputFile != "",
|
||
|
|
"Input file does not have extension .zst")) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
auto outputFd = openOutputFile(options, outputFile, state);
|
||
|
|
if (outputFd == nullptr) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
auto closeOutputGuard = makeScopeGuard([&] { std::fclose(outputFd); });
|
||
|
|
// (de)compress the file
|
||
|
|
handleOneInput(options, input, inputFd, outputFile, outputFd, state);
|
||
|
|
if (state.errorHolder.hasError()) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
// Delete the input file if necessary
|
||
|
|
if (!options.keepSource) {
|
||
|
|
// Be sure that we are done and have written everything before we delete
|
||
|
|
if (!state.errorHolder.check(std::fclose(inputFd) == 0,
|
||
|
|
"Failed to close input file")) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
closeInputGuard.dismiss();
|
||
|
|
if (!state.errorHolder.check(std::fclose(outputFd) == 0,
|
||
|
|
"Failed to close output file")) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
closeOutputGuard.dismiss();
|
||
|
|
if (std::remove(input.c_str()) != 0) {
|
||
|
|
state.errorHolder.setError("Failed to remove input file");
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Returns 1 if any of the files failed to (de)compress.
|
||
|
|
return returnCode;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Construct a `ZSTD_inBuffer` that points to the data in `buffer`.
|
||
|
|
static ZSTD_inBuffer makeZstdInBuffer(const Buffer& buffer) {
|
||
|
|
return ZSTD_inBuffer{buffer.data(), buffer.size(), 0};
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Advance `buffer` and `inBuffer` by the amount of data read, as indicated by
|
||
|
|
* `inBuffer.pos`.
|
||
|
|
*/
|
||
|
|
void advance(Buffer& buffer, ZSTD_inBuffer& inBuffer) {
|
||
|
|
auto pos = inBuffer.pos;
|
||
|
|
inBuffer.src = static_cast<const unsigned char*>(inBuffer.src) + pos;
|
||
|
|
inBuffer.size -= pos;
|
||
|
|
inBuffer.pos = 0;
|
||
|
|
return buffer.advance(pos);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Construct a `ZSTD_outBuffer` that points to the data in `buffer`.
|
||
|
|
static ZSTD_outBuffer makeZstdOutBuffer(Buffer& buffer) {
|
||
|
|
return ZSTD_outBuffer{buffer.data(), buffer.size(), 0};
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Split `buffer` and advance `outBuffer` by the amount of data written, as
|
||
|
|
* indicated by `outBuffer.pos`.
|
||
|
|
*/
|
||
|
|
Buffer split(Buffer& buffer, ZSTD_outBuffer& outBuffer) {
|
||
|
|
auto pos = outBuffer.pos;
|
||
|
|
outBuffer.dst = static_cast<unsigned char*>(outBuffer.dst) + pos;
|
||
|
|
outBuffer.size -= pos;
|
||
|
|
outBuffer.pos = 0;
|
||
|
|
return buffer.splitAt(pos);
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Stream chunks of input from `in`, compress it, and stream it out to `out`.
|
||
|
|
*
|
||
|
|
* @param state The shared state
|
||
|
|
* @param in Queue that we `pop()` input buffers from
|
||
|
|
* @param out Queue that we `push()` compressed output buffers to
|
||
|
|
* @param maxInputSize An upper bound on the size of the input
|
||
|
|
*/
|
||
|
|
static void compress(
|
||
|
|
SharedState& state,
|
||
|
|
std::shared_ptr<BufferWorkQueue> in,
|
||
|
|
std::shared_ptr<BufferWorkQueue> out,
|
||
|
|
size_t maxInputSize) {
|
||
|
|
auto& errorHolder = state.errorHolder;
|
||
|
|
auto guard = makeScopeGuard([&] { out->finish(); });
|
||
|
|
// Initialize the CCtx
|
||
|
|
auto ctx = state.cStreamPool->get();
|
||
|
|
if (!errorHolder.check(ctx != nullptr, "Failed to allocate ZSTD_CStream")) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
{
|
||
|
|
auto err = ZSTD_CCtx_reset(ctx.get(), ZSTD_reset_session_only);
|
||
|
|
if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Allocate space for the result
|
||
|
|
auto outBuffer = Buffer(ZSTD_compressBound(maxInputSize));
|
||
|
|
auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
|
||
|
|
{
|
||
|
|
Buffer inBuffer;
|
||
|
|
// Read a buffer in from the input queue
|
||
|
|
while (in->pop(inBuffer) && !errorHolder.hasError()) {
|
||
|
|
auto zstdInBuffer = makeZstdInBuffer(inBuffer);
|
||
|
|
// Compress the whole buffer and send it to the output queue
|
||
|
|
while (!inBuffer.empty() && !errorHolder.hasError()) {
|
||
|
|
if (!errorHolder.check(
|
||
|
|
!outBuffer.empty(), "ZSTD_compressBound() was too small")) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
// Compress
|
||
|
|
auto err =
|
||
|
|
ZSTD_compressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
|
||
|
|
if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
// Split the compressed data off outBuffer and pass to the output queue
|
||
|
|
out->push(split(outBuffer, zstdOutBuffer));
|
||
|
|
// Forget about the data we already compressed
|
||
|
|
advance(inBuffer, zstdInBuffer);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Write the epilog
|
||
|
|
size_t bytesLeft;
|
||
|
|
do {
|
||
|
|
if (!errorHolder.check(
|
||
|
|
!outBuffer.empty(), "ZSTD_compressBound() was too small")) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
bytesLeft = ZSTD_endStream(ctx.get(), &zstdOutBuffer);
|
||
|
|
if (!errorHolder.check(
|
||
|
|
!ZSTD_isError(bytesLeft), ZSTD_getErrorName(bytesLeft))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
out->push(split(outBuffer, zstdOutBuffer));
|
||
|
|
} while (bytesLeft != 0 && !errorHolder.hasError());
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Calculates how large each independently compressed frame should be.
|
||
|
|
*
|
||
|
|
* @param size The size of the source if known, 0 otherwise
|
||
|
|
* @param numThreads The number of threads available to run compression jobs on
|
||
|
|
* @param params The zstd parameters to be used for compression
|
||
|
|
*/
|
||
|
|
static size_t calculateStep(
|
||
|
|
std::uintmax_t size,
|
||
|
|
size_t numThreads,
|
||
|
|
const ZSTD_parameters ¶ms) {
|
||
|
|
(void)size;
|
||
|
|
(void)numThreads;
|
||
|
|
// Not validated to work correctly for window logs > 23.
|
||
|
|
// It will definitely fail if windowLog + 2 is >= 4GB because
|
||
|
|
// the skippable frame can only store sizes up to 4GB.
|
||
|
|
assert(params.cParams.windowLog <= 23);
|
||
|
|
return size_t{1} << (params.cParams.windowLog + 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
enum class FileStatus { Continue, Done, Error };
|
||
|
|
/// Determines the status of the file descriptor `fd`.
|
||
|
|
FileStatus fileStatus(FILE* fd) {
|
||
|
|
if (std::feof(fd)) {
|
||
|
|
return FileStatus::Done;
|
||
|
|
} else if (std::ferror(fd)) {
|
||
|
|
return FileStatus::Error;
|
||
|
|
}
|
||
|
|
return FileStatus::Continue;
|
||
|
|
}
|
||
|
|
} // anonymous namespace
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Reads `size` data in chunks of `chunkSize` and puts it into `queue`.
|
||
|
|
* Will read less if an error or EOF occurs.
|
||
|
|
* Returns the status of the file after all of the reads have occurred.
|
||
|
|
*/
|
||
|
|
static FileStatus
|
||
|
|
readData(BufferWorkQueue& queue, size_t chunkSize, size_t size, FILE* fd,
|
||
|
|
std::uint64_t *totalBytesRead) {
|
||
|
|
Buffer buffer(size);
|
||
|
|
while (!buffer.empty()) {
|
||
|
|
auto bytesRead =
|
||
|
|
std::fread(buffer.data(), 1, std::min(chunkSize, buffer.size()), fd);
|
||
|
|
*totalBytesRead += bytesRead;
|
||
|
|
queue.push(buffer.splitAt(bytesRead));
|
||
|
|
auto status = fileStatus(fd);
|
||
|
|
if (status != FileStatus::Continue) {
|
||
|
|
return status;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return FileStatus::Continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::uint64_t asyncCompressChunks(
|
||
|
|
SharedState& state,
|
||
|
|
WorkQueue<std::shared_ptr<BufferWorkQueue>>& chunks,
|
||
|
|
ThreadPool& executor,
|
||
|
|
FILE* fd,
|
||
|
|
std::uintmax_t size,
|
||
|
|
size_t numThreads,
|
||
|
|
ZSTD_parameters params) {
|
||
|
|
auto chunksGuard = makeScopeGuard([&] { chunks.finish(); });
|
||
|
|
std::uint64_t bytesRead = 0;
|
||
|
|
|
||
|
|
// Break the input up into chunks of size `step` and compress each chunk
|
||
|
|
// independently.
|
||
|
|
size_t step = calculateStep(size, numThreads, params);
|
||
|
|
state.log(kLogDebug, "Chosen frame size: %zu\n", step);
|
||
|
|
auto status = FileStatus::Continue;
|
||
|
|
while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
|
||
|
|
// Make a new input queue that we will put the chunk's input data into.
|
||
|
|
auto in = std::make_shared<BufferWorkQueue>();
|
||
|
|
auto inGuard = makeScopeGuard([&] { in->finish(); });
|
||
|
|
// Make a new output queue that compress will put the compressed data into.
|
||
|
|
auto out = std::make_shared<BufferWorkQueue>();
|
||
|
|
// Start compression in the thread pool
|
||
|
|
executor.add([&state, in, out, step] {
|
||
|
|
return compress(
|
||
|
|
state, std::move(in), std::move(out), step);
|
||
|
|
});
|
||
|
|
// Pass the output queue to the writer thread.
|
||
|
|
chunks.push(std::move(out));
|
||
|
|
state.log(kLogVerbose, "%s\n", "Starting a new frame");
|
||
|
|
// Fill the input queue for the compression job we just started
|
||
|
|
status = readData(*in, ZSTD_CStreamInSize(), step, fd, &bytesRead);
|
||
|
|
}
|
||
|
|
state.errorHolder.check(status != FileStatus::Error, "Error reading input");
|
||
|
|
return bytesRead;
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Decompress a frame, whose data is streamed into `in`, and stream the output
|
||
|
|
* to `out`.
|
||
|
|
*
|
||
|
|
* @param state The shared state
|
||
|
|
* @param in Queue that we `pop()` input buffers from. It contains
|
||
|
|
* exactly one compressed frame.
|
||
|
|
* @param out Queue that we `push()` decompressed output buffers to
|
||
|
|
*/
|
||
|
|
static void decompress(
|
||
|
|
SharedState& state,
|
||
|
|
std::shared_ptr<BufferWorkQueue> in,
|
||
|
|
std::shared_ptr<BufferWorkQueue> out) {
|
||
|
|
auto& errorHolder = state.errorHolder;
|
||
|
|
auto guard = makeScopeGuard([&] { out->finish(); });
|
||
|
|
// Initialize the DCtx
|
||
|
|
auto ctx = state.dStreamPool->get();
|
||
|
|
if (!errorHolder.check(ctx != nullptr, "Failed to allocate ZSTD_DStream")) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
{
|
||
|
|
auto err = ZSTD_DCtx_reset(ctx.get(), ZSTD_reset_session_only);
|
||
|
|
if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
const size_t outSize = ZSTD_DStreamOutSize();
|
||
|
|
Buffer inBuffer;
|
||
|
|
size_t returnCode = 0;
|
||
|
|
// Read a buffer in from the input queue
|
||
|
|
while (in->pop(inBuffer) && !errorHolder.hasError()) {
|
||
|
|
auto zstdInBuffer = makeZstdInBuffer(inBuffer);
|
||
|
|
// Decompress the whole buffer and send it to the output queue
|
||
|
|
while (!inBuffer.empty() && !errorHolder.hasError()) {
|
||
|
|
// Allocate a buffer with at least outSize bytes.
|
||
|
|
Buffer outBuffer(outSize);
|
||
|
|
auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
|
||
|
|
// Decompress
|
||
|
|
returnCode =
|
||
|
|
ZSTD_decompressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
|
||
|
|
if (!errorHolder.check(
|
||
|
|
!ZSTD_isError(returnCode), ZSTD_getErrorName(returnCode))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
// Pass the buffer with the decompressed data to the output queue
|
||
|
|
out->push(split(outBuffer, zstdOutBuffer));
|
||
|
|
// Advance past the input we already read
|
||
|
|
advance(inBuffer, zstdInBuffer);
|
||
|
|
if (returnCode == 0) {
|
||
|
|
// The frame is over, prepare to (maybe) start a new frame
|
||
|
|
ZSTD_initDStream(ctx.get());
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (!errorHolder.check(returnCode <= 1, "Incomplete block")) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
// We've given ZSTD_decompressStream all of our data, but there may still
|
||
|
|
// be data to read.
|
||
|
|
while (returnCode == 1) {
|
||
|
|
// Allocate a buffer with at least outSize bytes.
|
||
|
|
Buffer outBuffer(outSize);
|
||
|
|
auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
|
||
|
|
// Pass in no input.
|
||
|
|
ZSTD_inBuffer zstdInBuffer{nullptr, 0, 0};
|
||
|
|
// Decompress
|
||
|
|
returnCode =
|
||
|
|
ZSTD_decompressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
|
||
|
|
if (!errorHolder.check(
|
||
|
|
!ZSTD_isError(returnCode), ZSTD_getErrorName(returnCode))) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
// Pass the buffer with the decompressed data to the output queue
|
||
|
|
out->push(split(outBuffer, zstdOutBuffer));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
std::uint64_t asyncDecompressFrames(
|
||
|
|
SharedState& state,
|
||
|
|
WorkQueue<std::shared_ptr<BufferWorkQueue>>& frames,
|
||
|
|
ThreadPool& executor,
|
||
|
|
FILE* fd) {
|
||
|
|
auto framesGuard = makeScopeGuard([&] { frames.finish(); });
|
||
|
|
std::uint64_t totalBytesRead = 0;
|
||
|
|
|
||
|
|
// Split the source up into its component frames.
|
||
|
|
// If we find our recognized skippable frame we know the next frames size
|
||
|
|
// which means that we can decompress each standard frame in independently.
|
||
|
|
// Otherwise, we will decompress using only one decompression task.
|
||
|
|
const size_t chunkSize = ZSTD_DStreamInSize();
|
||
|
|
auto status = FileStatus::Continue;
|
||
|
|
while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
|
||
|
|
// Make a new input queue that we will put the frames's bytes into.
|
||
|
|
auto in = std::make_shared<BufferWorkQueue>();
|
||
|
|
auto inGuard = makeScopeGuard([&] { in->finish(); });
|
||
|
|
// Make a output queue that decompress will put the decompressed data into
|
||
|
|
auto out = std::make_shared<BufferWorkQueue>();
|
||
|
|
|
||
|
|
size_t frameSize;
|
||
|
|
{
|
||
|
|
// Calculate the size of the next frame.
|
||
|
|
// frameSize is 0 if the frame info can't be decoded.
|
||
|
|
Buffer buffer(SkippableFrame::kSize);
|
||
|
|
auto bytesRead = std::fread(buffer.data(), 1, buffer.size(), fd);
|
||
|
|
totalBytesRead += bytesRead;
|
||
|
|
status = fileStatus(fd);
|
||
|
|
if (bytesRead == 0 && status != FileStatus::Continue) {
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
buffer.subtract(buffer.size() - bytesRead);
|
||
|
|
frameSize = SkippableFrame::tryRead(buffer.range());
|
||
|
|
in->push(std::move(buffer));
|
||
|
|
}
|
||
|
|
if (frameSize == 0) {
|
||
|
|
// We hit a non SkippableFrame, so this will be the last job.
|
||
|
|
// Make sure that we don't use too much memory
|
||
|
|
in->setMaxSize(64);
|
||
|
|
out->setMaxSize(64);
|
||
|
|
}
|
||
|
|
// Start decompression in the thread pool
|
||
|
|
executor.add([&state, in, out] {
|
||
|
|
return decompress(state, std::move(in), std::move(out));
|
||
|
|
});
|
||
|
|
// Pass the output queue to the writer thread
|
||
|
|
frames.push(std::move(out));
|
||
|
|
if (frameSize == 0) {
|
||
|
|
// We hit a non SkippableFrame ==> not compressed by pzstd or corrupted
|
||
|
|
// Pass the rest of the source to this decompression task
|
||
|
|
state.log(kLogVerbose, "%s\n",
|
||
|
|
"Input not in pzstd format, falling back to serial decompression");
|
||
|
|
while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
|
||
|
|
status = readData(*in, chunkSize, chunkSize, fd, &totalBytesRead);
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
state.log(kLogVerbose, "Decompressing a frame of size %zu", frameSize);
|
||
|
|
// Fill the input queue for the decompression job we just started
|
||
|
|
status = readData(*in, chunkSize, frameSize, fd, &totalBytesRead);
|
||
|
|
}
|
||
|
|
state.errorHolder.check(status != FileStatus::Error, "Error reading input");
|
||
|
|
return totalBytesRead;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Write `data` to `fd`, returns true iff success.
|
||
|
|
static bool writeData(ByteRange data, FILE* fd) {
|
||
|
|
while (!data.empty()) {
|
||
|
|
data.advance(std::fwrite(data.begin(), 1, data.size(), fd));
|
||
|
|
if (std::ferror(fd)) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::uint64_t writeFile(
|
||
|
|
SharedState& state,
|
||
|
|
WorkQueue<std::shared_ptr<BufferWorkQueue>>& outs,
|
||
|
|
FILE* outputFd,
|
||
|
|
bool decompress) {
|
||
|
|
auto& errorHolder = state.errorHolder;
|
||
|
|
auto lineClearGuard = makeScopeGuard([&state] {
|
||
|
|
state.log.clear(kLogInfo);
|
||
|
|
});
|
||
|
|
std::uint64_t bytesWritten = 0;
|
||
|
|
std::shared_ptr<BufferWorkQueue> out;
|
||
|
|
// Grab the output queue for each decompression job (in order).
|
||
|
|
while (outs.pop(out)) {
|
||
|
|
if (errorHolder.hasError()) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if (!decompress) {
|
||
|
|
// If we are compressing and want to write skippable frames we can't
|
||
|
|
// start writing before compression is done because we need to know the
|
||
|
|
// compressed size.
|
||
|
|
// Wait for the compressed size to be available and write skippable frame
|
||
|
|
assert(uint64_t(out->size()) < uint64_t(1) << 32);
|
||
|
|
SkippableFrame frame(uint32_t(out->size()));
|
||
|
|
if (!writeData(frame.data(), outputFd)) {
|
||
|
|
errorHolder.setError("Failed to write output");
|
||
|
|
return bytesWritten;
|
||
|
|
}
|
||
|
|
bytesWritten += frame.kSize;
|
||
|
|
}
|
||
|
|
// For each chunk of the frame: Pop it from the queue and write it
|
||
|
|
Buffer buffer;
|
||
|
|
while (out->pop(buffer) && !errorHolder.hasError()) {
|
||
|
|
if (!writeData(buffer.range(), outputFd)) {
|
||
|
|
errorHolder.setError("Failed to write output");
|
||
|
|
return bytesWritten;
|
||
|
|
}
|
||
|
|
bytesWritten += buffer.size();
|
||
|
|
state.log.update(kLogInfo, "Written: %u MB ",
|
||
|
|
static_cast<std::uint32_t>(bytesWritten >> 20));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return bytesWritten;
|
||
|
|
}
|
||
|
|
}
|