1 /* 2 * Copyright (c) Martin Liska, SUSE, Facebook, Inc. 3 * All rights reserved. 4 * 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. 9 */ 10 11 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() 17 #include <pthread.h> 18 19 typedef struct compress_args 20 { 21 const char *fname; 22 char *outName; 23 int cLevel; 24 #if defined(ZSTD_STATIC_LINKING_ONLY) 25 ZSTD_threadPool *pool; 26 #endif 27 } compress_args_t; 28 29 static void *compressFile_orDie(void *data) 30 { 31 const int nbThreads = 16; 32 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. 41 */ 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); 46 47 /* Create the context. */ 48 ZSTD_CCtx* const cctx = ZSTD_createCCtx(); 49 CHECK(cctx != NULL, "ZSTD_createCCtx() failed!"); 50 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!"); 54 #endif 55 56 /* Set any parameters you want. 57 * Here we set the compression level, and enable the checksum. 58 */ 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); 62 63 /* This loop reads from the input file, compresses that entire chunk, 64 * and writes all output produced to the output file. 65 */ 66 size_t const toRead = buffInSize; 67 for (;;) { 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. 74 */ 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 79 * output. 80 */ 81 ZSTD_inBuffer input = { buffIn, read, 0 }; 82 int finished; 83 do { 84 /* Compress into the output buffer and write all of the output to 85 * the file so we can reuse the buffer next iteration. 86 */ 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. 94 */ 95 finished = lastChunk ? (remaining == 0) : (input.pos == input.size); 96 } while (!finished); 97 CHECK(input.pos == input.size, 98 "Impossible: zstd only returns 0 when the input is completely consumed!"); 99 100 if (lastChunk) { 101 break; 102 } 103 } 104 105 fprintf (stderr, "Finishing compression of %s\n", args->outName); 106 107 ZSTD_freeCCtx(cctx); 108 fclose_orDie(fout); 109 fclose_orDie(fin); 110 free(buffIn); 111 free(buffOut); 112 free(args->outName); 113 114 return NULL; 115 } 116 117 118 static char* createOutFilename_orDie(const char* filename) 119 { 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; 127 } 128 129 int main(int argc, const char** argv) 130 { 131 const char* const exeName = argv[0]; 132 133 if (argc<=3) { 134 printf("wrong arguments\n"); 135 printf("usage:\n"); 136 printf("%s POOL_SIZE LEVEL FILES\n", exeName); 137 return 1; 138 } 139 140 int pool_size = atoi (argv[1]); 141 CHECK(pool_size != 0, "can't parse POOL_SIZE!"); 142 143 int level = atoi (argv[2]); 144 CHECK(level != 0, "can't parse LEVEL!"); 145 146 argc -= 3; 147 argv += 3; 148 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); 153 #else 154 fprintf (stderr, "All threads use its own thread pool\n"); 155 #endif 156 157 pthread_t *threads = malloc_orDie(argc * sizeof(pthread_t)); 158 compress_args_t *args = malloc_orDie(argc * sizeof(compress_args_t)); 159 160 for (unsigned i = 0; i < argc; i++) 161 { 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) 166 args[i].pool = pool; 167 #endif 168 169 pthread_create (&threads[i], NULL, compressFile_orDie, &args[i]); 170 } 171 172 for (unsigned i = 0; i < argc; i++) 173 pthread_join (threads[i], NULL); 174 175 #if defined(ZSTD_STATIC_LINKING_ONLY) 176 ZSTD_freeThreadPool (pool); 177 #endif 178 179 return 0; 180 } 181