2 * Copyright (c) Martin Liska, SUSE, 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).
8 * You may select, at your option, one of the above-listed licenses.
12 #include <stdio.h> // printf
13 #include <stdlib.h> // free
14 #include <string.h> // memset, strcat, strlen
15 #include <zstd.h> // presumes zstd library is installed
16 #include "common.h" // Helper functions, CHECK(), and CHECK_ZSTD()
19 typedef struct compress_args
24 #if defined(ZSTD_STATIC_LINKING_ONLY)
25 ZSTD_threadPool *pool;
29 static void *compressFile_orDie(void *data)
31 const int nbThreads = 16;
33 compress_args_t *args = (compress_args_t *)data;
34 fprintf (stderr, "Starting compression of %s with level %d, using %d threads\n", args->fname, args->cLevel, nbThreads);
35 /* Open the input and output files. */
36 FILE* const fin = fopen_orDie(args->fname, "rb");
37 FILE* const fout = fopen_orDie(args->outName, "wb");
38 /* Create the input and output buffers.
39 * They may be any size, but we recommend using these functions to size them.
40 * Performance will only suffer significantly for very tiny buffers.
42 size_t const buffInSize = ZSTD_CStreamInSize();
43 void* const buffIn = malloc_orDie(buffInSize);
44 size_t const buffOutSize = ZSTD_CStreamOutSize();
45 void* const buffOut = malloc_orDie(buffOutSize);
47 /* Create the context. */
48 ZSTD_CCtx* const cctx = ZSTD_createCCtx();
49 CHECK(cctx != NULL, "ZSTD_createCCtx() failed!");
51 #if defined(ZSTD_STATIC_LINKING_ONLY)
52 size_t r = ZSTD_CCtx_refThreadPool(cctx, args->pool);
53 CHECK(r == 0, "ZSTD_CCtx_refThreadPool failed!");
56 /* Set any parameters you want.
57 * Here we set the compression level, and enable the checksum.
59 CHECK_ZSTD( ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, args->cLevel) );
60 CHECK_ZSTD( ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 1) );
61 ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, nbThreads);
63 /* This loop reads from the input file, compresses that entire chunk,
64 * and writes all output produced to the output file.
66 size_t const toRead = buffInSize;
68 size_t read = fread_orDie(buffIn, toRead, fin);
69 /* Select the flush mode.
70 * If the read may not be finished (read == toRead) we use
71 * ZSTD_e_continue. If this is the last chunk, we use ZSTD_e_end.
72 * Zstd optimizes the case where the first flush mode is ZSTD_e_end,
73 * since it knows it is compressing the entire source in one pass.
75 int const lastChunk = (read < toRead);
76 ZSTD_EndDirective const mode = lastChunk ? ZSTD_e_end : ZSTD_e_continue;
77 /* Set the input buffer to what we just read.
78 * We compress until the input buffer is empty, each time flushing the
81 ZSTD_inBuffer input = { buffIn, read, 0 };
84 /* Compress into the output buffer and write all of the output to
85 * the file so we can reuse the buffer next iteration.
87 ZSTD_outBuffer output = { buffOut, buffOutSize, 0 };
88 size_t const remaining = ZSTD_compressStream2(cctx, &output , &input, mode);
89 CHECK_ZSTD(remaining);
90 fwrite_orDie(buffOut, output.pos, fout);
91 /* If we're on the last chunk we're finished when zstd returns 0,
92 * which means its consumed all the input AND finished the frame.
93 * Otherwise, we're finished when we've consumed all the input.
95 finished = lastChunk ? (remaining == 0) : (input.pos == input.size);
97 CHECK(input.pos == input.size,
98 "Impossible: zstd only returns 0 when the input is completely consumed!");
105 fprintf (stderr, "Finishing compression of %s\n", args->outName);
118 static char* createOutFilename_orDie(const char* filename)
120 size_t const inL = strlen(filename);
121 size_t const outL = inL + 5;
122 void* const outSpace = malloc_orDie(outL);
123 memset(outSpace, 0, outL);
124 strcat(outSpace, filename);
125 strcat(outSpace, ".zst");
126 return (char*)outSpace;
129 int main(int argc, const char** argv)
131 const char* const exeName = argv[0];
134 printf("wrong arguments\n");
136 printf("%s POOL_SIZE LEVEL FILES\n", exeName);
140 int pool_size = atoi (argv[1]);
141 CHECK(pool_size != 0, "can't parse POOL_SIZE!");
143 int level = atoi (argv[2]);
144 CHECK(level != 0, "can't parse LEVEL!");
149 #if defined(ZSTD_STATIC_LINKING_ONLY)
150 ZSTD_threadPool *pool = ZSTD_createThreadPool (pool_size);
151 CHECK(pool != NULL, "ZSTD_createThreadPool() failed!");
152 fprintf (stderr, "Using shared thread pool of size %d\n", pool_size);
154 fprintf (stderr, "All threads use its own thread pool\n");
157 pthread_t *threads = malloc_orDie(argc * sizeof(pthread_t));
158 compress_args_t *args = malloc_orDie(argc * sizeof(compress_args_t));
160 for (unsigned i = 0; i < argc; i++)
162 args[i].fname = argv[i];
163 args[i].outName = createOutFilename_orDie(args[i].fname);
164 args[i].cLevel = level;
165 #if defined(ZSTD_STATIC_LINKING_ONLY)
169 pthread_create (&threads[i], NULL, compressFile_orDie, &args[i]);
172 for (unsigned i = 0; i < argc; i++)
173 pthread_join (threads[i], NULL);
175 #if defined(ZSTD_STATIC_LINKING_ONLY)
176 ZSTD_freeThreadPool (pool);