| /* |
| * Copyright (c) 2016-present, Facebook, Inc. |
| * 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/Range.h" |
| #include "utils/ScopeGuard.h" |
| #include "utils/ThreadPool.h" |
| #include "utils/WorkQueue.h" |
| |
| #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(INFO, "%-20s :%6.2f%% (%6" PRIu64 " => %6" PRIu64 |
| " bytes, %s)\n", |
| inputFileName.c_str(), ratio * 100, bytesRead, bytesWritten, |
| outputFileName.c_str()); |
| } else { |
| state.log(INFO, "%-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(INFO)) { |
| state.errorHolder.setError("Output file exists"); |
| return nullptr; |
| } |
| state.log( |
| INFO, |
| "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(ERROR, "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_resetCStream(ctx.get(), 0); |
| 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; |
| 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(DEBUG, "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(VERBOSE, "%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_resetDStream(ctx.get()); |
| 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(VERBOSE, "%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(VERBOSE, "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(INFO); |
| }); |
| 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 |
| SkippableFrame frame(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(INFO, "Written: %u MB ", |
| static_cast<std::uint32_t>(bytesWritten >> 20)); |
| } |
| } |
| return bytesWritten; |
| } |
| } |
