2 * Copyright (c) Meta Platforms, Inc. and affiliates.
5 * This source code is licensed under both the BSD-style license (found in the
6 * LICENSE file in the root directory of this source tree) and the GPLv2 (found
7 * in the COPYING file in the root directory of this source tree).
9 #include "platform.h" /* Large Files support, SET_BINARY_MODE */
11 #include "SkippableFrame.h"
12 #include "utils/FileSystem.h"
13 #include "utils/Portability.h"
14 #include "utils/Range.h"
15 #include "utils/ScopeGuard.h"
16 #include "utils/ThreadPool.h"
17 #include "utils/WorkQueue.h"
32 const std::string nullOutput = "nul";
34 const std::string nullOutput = "/dev/null";
40 static std::uintmax_t fileSizeOrZero(const std::string &file) {
45 auto size = file_size(file, ec);
52 static std::uint64_t handleOneInput(const Options &options,
53 const std::string &inputFile,
55 const std::string &outputFile,
58 auto inputSize = fileSizeOrZero(inputFile);
59 // WorkQueue outlives ThreadPool so in the case of error we are certain
60 // we don't accidentally try to call push() on it after it is destroyed
61 WorkQueue<std::shared_ptr<BufferWorkQueue>> outs{options.numThreads + 1};
62 std::uint64_t bytesRead;
63 std::uint64_t bytesWritten;
65 // Initialize the (de)compression thread pool with numThreads
66 ThreadPool executor(options.numThreads);
67 // Run the reader thread on an extra thread
68 ThreadPool readExecutor(1);
69 if (!options.decompress) {
70 // Add a job that reads the input and starts all the compression jobs
72 [&state, &outs, &executor, inputFd, inputSize, &options, &bytesRead] {
73 bytesRead = asyncCompressChunks(
80 options.determineParameters());
83 bytesWritten = writeFile(state, outs, outputFd, options.decompress);
85 // Add a job that reads the input and starts all the decompression jobs
86 readExecutor.add([&state, &outs, &executor, inputFd, &bytesRead] {
87 bytesRead = asyncDecompressFrames(state, outs, executor, inputFd);
90 bytesWritten = writeFile(state, outs, outputFd, options.decompress);
93 if (!state.errorHolder.hasError()) {
94 std::string inputFileName = inputFile == "-" ? "stdin" : inputFile;
95 std::string outputFileName = outputFile == "-" ? "stdout" : outputFile;
96 if (!options.decompress) {
97 double ratio = static_cast<double>(bytesWritten) /
98 static_cast<double>(bytesRead + !bytesRead);
99 state.log(kLogInfo, "%-20s :%6.2f%% (%6" PRIu64 " => %6" PRIu64
101 inputFileName.c_str(), ratio * 100, bytesRead, bytesWritten,
102 outputFileName.c_str());
104 state.log(kLogInfo, "%-20s: %" PRIu64 " bytes \n",
105 inputFileName.c_str(),bytesWritten);
111 static FILE *openInputFile(const std::string &inputFile,
112 ErrorHolder &errorHolder) {
113 if (inputFile == "-") {
114 SET_BINARY_MODE(stdin);
117 // Check if input file is a directory
120 if (is_directory(inputFile, ec)) {
121 errorHolder.setError("Output file is a directory -- ignored");
125 auto inputFd = std::fopen(inputFile.c_str(), "rb");
126 if (!errorHolder.check(inputFd != nullptr, "Failed to open input file")) {
132 static FILE *openOutputFile(const Options &options,
133 const std::string &outputFile,
134 SharedState& state) {
135 if (outputFile == "-") {
136 SET_BINARY_MODE(stdout);
139 // Check if the output file exists and then open it
140 if (!options.overwrite && outputFile != nullOutput) {
141 auto outputFd = std::fopen(outputFile.c_str(), "rb");
142 if (outputFd != nullptr) {
143 std::fclose(outputFd);
144 if (!state.log.logsAt(kLogInfo)) {
145 state.errorHolder.setError("Output file exists");
150 "pzstd: %s already exists; do you wish to overwrite (y/n) ? ",
153 if (c != 'y' && c != 'Y') {
154 state.errorHolder.setError("Not overwritten");
159 auto outputFd = std::fopen(outputFile.c_str(), "wb");
160 if (!state.errorHolder.check(
161 outputFd != nullptr, "Failed to open output file")) {
167 int pzstdMain(const Options &options) {
169 SharedState state(options);
170 for (const auto& input : options.inputFiles) {
171 // Setup the shared state
172 auto printErrorGuard = makeScopeGuard([&] {
173 if (state.errorHolder.hasError()) {
175 state.log(kLogError, "pzstd: %s: %s.\n", input.c_str(),
176 state.errorHolder.getError().c_str());
179 // Open the input file
180 auto inputFd = openInputFile(input, state.errorHolder);
181 if (inputFd == nullptr) {
184 auto closeInputGuard = makeScopeGuard([&] { std::fclose(inputFd); });
185 // Open the output file
186 auto outputFile = options.getOutputFile(input);
187 if (!state.errorHolder.check(outputFile != "",
188 "Input file does not have extension .zst")) {
191 auto outputFd = openOutputFile(options, outputFile, state);
192 if (outputFd == nullptr) {
195 auto closeOutputGuard = makeScopeGuard([&] { std::fclose(outputFd); });
196 // (de)compress the file
197 handleOneInput(options, input, inputFd, outputFile, outputFd, state);
198 if (state.errorHolder.hasError()) {
201 // Delete the input file if necessary
202 if (!options.keepSource) {
203 // Be sure that we are done and have written everything before we delete
204 if (!state.errorHolder.check(std::fclose(inputFd) == 0,
205 "Failed to close input file")) {
208 closeInputGuard.dismiss();
209 if (!state.errorHolder.check(std::fclose(outputFd) == 0,
210 "Failed to close output file")) {
213 closeOutputGuard.dismiss();
214 if (std::remove(input.c_str()) != 0) {
215 state.errorHolder.setError("Failed to remove input file");
220 // Returns 1 if any of the files failed to (de)compress.
224 /// Construct a `ZSTD_inBuffer` that points to the data in `buffer`.
225 static ZSTD_inBuffer makeZstdInBuffer(const Buffer& buffer) {
226 return ZSTD_inBuffer{buffer.data(), buffer.size(), 0};
230 * Advance `buffer` and `inBuffer` by the amount of data read, as indicated by
233 void advance(Buffer& buffer, ZSTD_inBuffer& inBuffer) {
234 auto pos = inBuffer.pos;
235 inBuffer.src = static_cast<const unsigned char*>(inBuffer.src) + pos;
236 inBuffer.size -= pos;
238 return buffer.advance(pos);
241 /// Construct a `ZSTD_outBuffer` that points to the data in `buffer`.
242 static ZSTD_outBuffer makeZstdOutBuffer(Buffer& buffer) {
243 return ZSTD_outBuffer{buffer.data(), buffer.size(), 0};
247 * Split `buffer` and advance `outBuffer` by the amount of data written, as
248 * indicated by `outBuffer.pos`.
250 Buffer split(Buffer& buffer, ZSTD_outBuffer& outBuffer) {
251 auto pos = outBuffer.pos;
252 outBuffer.dst = static_cast<unsigned char*>(outBuffer.dst) + pos;
253 outBuffer.size -= pos;
255 return buffer.splitAt(pos);
259 * Stream chunks of input from `in`, compress it, and stream it out to `out`.
261 * @param state The shared state
262 * @param in Queue that we `pop()` input buffers from
263 * @param out Queue that we `push()` compressed output buffers to
264 * @param maxInputSize An upper bound on the size of the input
266 static void compress(
268 std::shared_ptr<BufferWorkQueue> in,
269 std::shared_ptr<BufferWorkQueue> out,
270 size_t maxInputSize) {
271 auto& errorHolder = state.errorHolder;
272 auto guard = makeScopeGuard([&] { out->finish(); });
273 // Initialize the CCtx
274 auto ctx = state.cStreamPool->get();
275 if (!errorHolder.check(ctx != nullptr, "Failed to allocate ZSTD_CStream")) {
279 auto err = ZSTD_CCtx_reset(ctx.get(), ZSTD_reset_session_only);
280 if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
285 // Allocate space for the result
286 auto outBuffer = Buffer(ZSTD_compressBound(maxInputSize));
287 auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
290 // Read a buffer in from the input queue
291 while (in->pop(inBuffer) && !errorHolder.hasError()) {
292 auto zstdInBuffer = makeZstdInBuffer(inBuffer);
293 // Compress the whole buffer and send it to the output queue
294 while (!inBuffer.empty() && !errorHolder.hasError()) {
295 if (!errorHolder.check(
296 !outBuffer.empty(), "ZSTD_compressBound() was too small")) {
301 ZSTD_compressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
302 if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
305 // Split the compressed data off outBuffer and pass to the output queue
306 out->push(split(outBuffer, zstdOutBuffer));
307 // Forget about the data we already compressed
308 advance(inBuffer, zstdInBuffer);
315 if (!errorHolder.check(
316 !outBuffer.empty(), "ZSTD_compressBound() was too small")) {
319 bytesLeft = ZSTD_endStream(ctx.get(), &zstdOutBuffer);
320 if (!errorHolder.check(
321 !ZSTD_isError(bytesLeft), ZSTD_getErrorName(bytesLeft))) {
324 out->push(split(outBuffer, zstdOutBuffer));
325 } while (bytesLeft != 0 && !errorHolder.hasError());
329 * Calculates how large each independently compressed frame should be.
331 * @param size The size of the source if known, 0 otherwise
332 * @param numThreads The number of threads available to run compression jobs on
333 * @param params The zstd parameters to be used for compression
335 static size_t calculateStep(
338 const ZSTD_parameters ¶ms) {
341 // Not validated to work correctly for window logs > 23.
342 // It will definitely fail if windowLog + 2 is >= 4GB because
343 // the skippable frame can only store sizes up to 4GB.
344 assert(params.cParams.windowLog <= 23);
345 return size_t{1} << (params.cParams.windowLog + 2);
349 enum class FileStatus { Continue, Done, Error };
350 /// Determines the status of the file descriptor `fd`.
351 FileStatus fileStatus(FILE* fd) {
353 return FileStatus::Done;
354 } else if (std::ferror(fd)) {
355 return FileStatus::Error;
357 return FileStatus::Continue;
359 } // anonymous namespace
362 * Reads `size` data in chunks of `chunkSize` and puts it into `queue`.
363 * Will read less if an error or EOF occurs.
364 * Returns the status of the file after all of the reads have occurred.
367 readData(BufferWorkQueue& queue, size_t chunkSize, size_t size, FILE* fd,
368 std::uint64_t *totalBytesRead) {
370 while (!buffer.empty()) {
372 std::fread(buffer.data(), 1, std::min(chunkSize, buffer.size()), fd);
373 *totalBytesRead += bytesRead;
374 queue.push(buffer.splitAt(bytesRead));
375 auto status = fileStatus(fd);
376 if (status != FileStatus::Continue) {
380 return FileStatus::Continue;
383 std::uint64_t asyncCompressChunks(
385 WorkQueue<std::shared_ptr<BufferWorkQueue>>& chunks,
386 ThreadPool& executor,
390 ZSTD_parameters params) {
391 auto chunksGuard = makeScopeGuard([&] { chunks.finish(); });
392 std::uint64_t bytesRead = 0;
394 // Break the input up into chunks of size `step` and compress each chunk
396 size_t step = calculateStep(size, numThreads, params);
397 state.log(kLogDebug, "Chosen frame size: %zu\n", step);
398 auto status = FileStatus::Continue;
399 while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
400 // Make a new input queue that we will put the chunk's input data into.
401 auto in = std::make_shared<BufferWorkQueue>();
402 auto inGuard = makeScopeGuard([&] { in->finish(); });
403 // Make a new output queue that compress will put the compressed data into.
404 auto out = std::make_shared<BufferWorkQueue>();
405 // Start compression in the thread pool
406 executor.add([&state, in, out, step] {
408 state, std::move(in), std::move(out), step);
410 // Pass the output queue to the writer thread.
411 chunks.push(std::move(out));
412 state.log(kLogVerbose, "%s\n", "Starting a new frame");
413 // Fill the input queue for the compression job we just started
414 status = readData(*in, ZSTD_CStreamInSize(), step, fd, &bytesRead);
416 state.errorHolder.check(status != FileStatus::Error, "Error reading input");
421 * Decompress a frame, whose data is streamed into `in`, and stream the output
424 * @param state The shared state
425 * @param in Queue that we `pop()` input buffers from. It contains
426 * exactly one compressed frame.
427 * @param out Queue that we `push()` decompressed output buffers to
429 static void decompress(
431 std::shared_ptr<BufferWorkQueue> in,
432 std::shared_ptr<BufferWorkQueue> out) {
433 auto& errorHolder = state.errorHolder;
434 auto guard = makeScopeGuard([&] { out->finish(); });
435 // Initialize the DCtx
436 auto ctx = state.dStreamPool->get();
437 if (!errorHolder.check(ctx != nullptr, "Failed to allocate ZSTD_DStream")) {
441 auto err = ZSTD_DCtx_reset(ctx.get(), ZSTD_reset_session_only);
442 if (!errorHolder.check(!ZSTD_isError(err), ZSTD_getErrorName(err))) {
447 const size_t outSize = ZSTD_DStreamOutSize();
449 size_t returnCode = 0;
450 // Read a buffer in from the input queue
451 while (in->pop(inBuffer) && !errorHolder.hasError()) {
452 auto zstdInBuffer = makeZstdInBuffer(inBuffer);
453 // Decompress the whole buffer and send it to the output queue
454 while (!inBuffer.empty() && !errorHolder.hasError()) {
455 // Allocate a buffer with at least outSize bytes.
456 Buffer outBuffer(outSize);
457 auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
460 ZSTD_decompressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
461 if (!errorHolder.check(
462 !ZSTD_isError(returnCode), ZSTD_getErrorName(returnCode))) {
465 // Pass the buffer with the decompressed data to the output queue
466 out->push(split(outBuffer, zstdOutBuffer));
467 // Advance past the input we already read
468 advance(inBuffer, zstdInBuffer);
469 if (returnCode == 0) {
470 // The frame is over, prepare to (maybe) start a new frame
471 ZSTD_initDStream(ctx.get());
475 if (!errorHolder.check(returnCode <= 1, "Incomplete block")) {
478 // We've given ZSTD_decompressStream all of our data, but there may still
480 while (returnCode == 1) {
481 // Allocate a buffer with at least outSize bytes.
482 Buffer outBuffer(outSize);
483 auto zstdOutBuffer = makeZstdOutBuffer(outBuffer);
485 ZSTD_inBuffer zstdInBuffer{nullptr, 0, 0};
488 ZSTD_decompressStream(ctx.get(), &zstdOutBuffer, &zstdInBuffer);
489 if (!errorHolder.check(
490 !ZSTD_isError(returnCode), ZSTD_getErrorName(returnCode))) {
493 // Pass the buffer with the decompressed data to the output queue
494 out->push(split(outBuffer, zstdOutBuffer));
498 std::uint64_t asyncDecompressFrames(
500 WorkQueue<std::shared_ptr<BufferWorkQueue>>& frames,
501 ThreadPool& executor,
503 auto framesGuard = makeScopeGuard([&] { frames.finish(); });
504 std::uint64_t totalBytesRead = 0;
506 // Split the source up into its component frames.
507 // If we find our recognized skippable frame we know the next frames size
508 // which means that we can decompress each standard frame in independently.
509 // Otherwise, we will decompress using only one decompression task.
510 const size_t chunkSize = ZSTD_DStreamInSize();
511 auto status = FileStatus::Continue;
512 while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
513 // Make a new input queue that we will put the frames's bytes into.
514 auto in = std::make_shared<BufferWorkQueue>();
515 auto inGuard = makeScopeGuard([&] { in->finish(); });
516 // Make a output queue that decompress will put the decompressed data into
517 auto out = std::make_shared<BufferWorkQueue>();
521 // Calculate the size of the next frame.
522 // frameSize is 0 if the frame info can't be decoded.
523 Buffer buffer(SkippableFrame::kSize);
524 auto bytesRead = std::fread(buffer.data(), 1, buffer.size(), fd);
525 totalBytesRead += bytesRead;
526 status = fileStatus(fd);
527 if (bytesRead == 0 && status != FileStatus::Continue) {
530 buffer.subtract(buffer.size() - bytesRead);
531 frameSize = SkippableFrame::tryRead(buffer.range());
532 in->push(std::move(buffer));
534 if (frameSize == 0) {
535 // We hit a non SkippableFrame, so this will be the last job.
536 // Make sure that we don't use too much memory
540 // Start decompression in the thread pool
541 executor.add([&state, in, out] {
542 return decompress(state, std::move(in), std::move(out));
544 // Pass the output queue to the writer thread
545 frames.push(std::move(out));
546 if (frameSize == 0) {
547 // We hit a non SkippableFrame ==> not compressed by pzstd or corrupted
548 // Pass the rest of the source to this decompression task
549 state.log(kLogVerbose, "%s\n",
550 "Input not in pzstd format, falling back to serial decompression");
551 while (status == FileStatus::Continue && !state.errorHolder.hasError()) {
552 status = readData(*in, chunkSize, chunkSize, fd, &totalBytesRead);
556 state.log(kLogVerbose, "Decompressing a frame of size %zu", frameSize);
557 // Fill the input queue for the decompression job we just started
558 status = readData(*in, chunkSize, frameSize, fd, &totalBytesRead);
560 state.errorHolder.check(status != FileStatus::Error, "Error reading input");
561 return totalBytesRead;
564 /// Write `data` to `fd`, returns true iff success.
565 static bool writeData(ByteRange data, FILE* fd) {
566 while (!data.empty()) {
567 data.advance(std::fwrite(data.begin(), 1, data.size(), fd));
568 if (std::ferror(fd)) {
575 std::uint64_t writeFile(
577 WorkQueue<std::shared_ptr<BufferWorkQueue>>& outs,
580 auto& errorHolder = state.errorHolder;
581 auto lineClearGuard = makeScopeGuard([&state] {
582 state.log.clear(kLogInfo);
584 std::uint64_t bytesWritten = 0;
585 std::shared_ptr<BufferWorkQueue> out;
586 // Grab the output queue for each decompression job (in order).
587 while (outs.pop(out)) {
588 if (errorHolder.hasError()) {
592 // If we are compressing and want to write skippable frames we can't
593 // start writing before compression is done because we need to know the
595 // Wait for the compressed size to be available and write skippable frame
596 assert(uint64_t(out->size()) < uint64_t(1) << 32);
597 SkippableFrame frame(uint32_t(out->size()));
598 if (!writeData(frame.data(), outputFd)) {
599 errorHolder.setError("Failed to write output");
602 bytesWritten += frame.kSize;
604 // For each chunk of the frame: Pop it from the queue and write it
606 while (out->pop(buffer) && !errorHolder.hasError()) {
607 if (!writeData(buffer.range(), outputFd)) {
608 errorHolder.setError("Failed to write output");
611 bytesWritten += buffer.size();
612 state.log.update(kLogInfo, "Written: %u MB ",
613 static_cast<std::uint32_t>(bytesWritten >> 20));