648db22b |
1 | /* |
2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
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 | */ |
9 | |
10 | #include <stdlib.h> // malloc, free, exit, atoi |
11 | #include <stdio.h> // fprintf, perror, feof, fopen, etc. |
12 | #include <string.h> // strlen, memset, strcat |
13 | #define ZSTD_STATIC_LINKING_ONLY |
14 | #include <zstd.h> // presumes zstd library is installed |
15 | #include <zstd_errors.h> |
16 | #if defined(WIN32) || defined(_WIN32) |
17 | # include <windows.h> |
18 | # define SLEEP(x) Sleep(x) |
19 | #else |
20 | # include <unistd.h> |
21 | # define SLEEP(x) usleep(x * 1000) |
22 | #endif |
23 | |
24 | #include "xxhash.h" |
25 | |
26 | #include "pool.h" // use zstd thread pool for demo |
27 | |
28 | #include "../zstd_seekable.h" |
29 | |
30 | static void* malloc_orDie(size_t size) |
31 | { |
32 | void* const buff = malloc(size); |
33 | if (buff) return buff; |
34 | /* error */ |
35 | perror("malloc:"); |
36 | exit(1); |
37 | } |
38 | |
39 | static FILE* fopen_orDie(const char *filename, const char *instruction) |
40 | { |
41 | FILE* const inFile = fopen(filename, instruction); |
42 | if (inFile) return inFile; |
43 | /* error */ |
44 | perror(filename); |
45 | exit(3); |
46 | } |
47 | |
48 | static size_t fread_orDie(void* buffer, size_t sizeToRead, FILE* file) |
49 | { |
50 | size_t const readSize = fread(buffer, 1, sizeToRead, file); |
51 | if (readSize == sizeToRead) return readSize; /* good */ |
52 | if (feof(file)) return readSize; /* good, reached end of file */ |
53 | /* error */ |
54 | perror("fread"); |
55 | exit(4); |
56 | } |
57 | |
58 | static size_t fwrite_orDie(const void* buffer, size_t sizeToWrite, FILE* file) |
59 | { |
60 | size_t const writtenSize = fwrite(buffer, 1, sizeToWrite, file); |
61 | if (writtenSize == sizeToWrite) return sizeToWrite; /* good */ |
62 | /* error */ |
63 | perror("fwrite"); |
64 | exit(5); |
65 | } |
66 | |
67 | static size_t fclose_orDie(FILE* file) |
68 | { |
69 | if (!fclose(file)) return 0; |
70 | /* error */ |
71 | perror("fclose"); |
72 | exit(6); |
73 | } |
74 | |
75 | static void fseek_orDie(FILE* file, long int offset, int origin) |
76 | { |
77 | if (!fseek(file, offset, origin)) { |
78 | if (!fflush(file)) return; |
79 | } |
80 | /* error */ |
81 | perror("fseek"); |
82 | exit(7); |
83 | } |
84 | |
85 | static long int ftell_orDie(FILE* file) |
86 | { |
87 | long int off = ftell(file); |
88 | if (off != -1) return off; |
89 | /* error */ |
90 | perror("ftell"); |
91 | exit(8); |
92 | } |
93 | |
94 | struct job { |
95 | const void* src; |
96 | size_t srcSize; |
97 | void* dst; |
98 | size_t dstSize; |
99 | |
100 | unsigned checksum; |
101 | |
102 | int compressionLevel; |
103 | int done; |
104 | }; |
105 | |
106 | static void compressFrame(void* opaque) |
107 | { |
108 | struct job* job = opaque; |
109 | |
110 | job->checksum = XXH64(job->src, job->srcSize, 0); |
111 | |
112 | size_t ret = ZSTD_compress(job->dst, job->dstSize, job->src, job->srcSize, job->compressionLevel); |
113 | if (ZSTD_isError(ret)) { |
114 | fprintf(stderr, "ZSTD_compress() error : %s \n", ZSTD_getErrorName(ret)); |
115 | exit(20); |
116 | } |
117 | |
118 | job->dstSize = ret; |
119 | job->done = 1; |
120 | } |
121 | |
122 | static void compressFile_orDie(const char* fname, const char* outName, int cLevel, unsigned frameSize, int nbThreads) |
123 | { |
124 | POOL_ctx* pool = POOL_create(nbThreads, nbThreads); |
125 | if (pool == NULL) { fprintf(stderr, "POOL_create() error \n"); exit(9); } |
126 | |
127 | FILE* const fin = fopen_orDie(fname, "rb"); |
128 | FILE* const fout = fopen_orDie(outName, "wb"); |
129 | |
130 | if (ZSTD_compressBound(frameSize) > 0xFFFFFFFFU) { fprintf(stderr, "Frame size too large \n"); exit(10); } |
131 | unsigned dstSize = ZSTD_compressBound(frameSize); |
132 | |
133 | |
134 | fseek_orDie(fin, 0, SEEK_END); |
135 | long int length = ftell_orDie(fin); |
136 | fseek_orDie(fin, 0, SEEK_SET); |
137 | |
138 | size_t numFrames = (length + frameSize - 1) / frameSize; |
139 | |
140 | struct job* jobs = malloc_orDie(sizeof(struct job) * numFrames); |
141 | |
142 | size_t i; |
143 | for(i = 0; i < numFrames; i++) { |
144 | void* in = malloc_orDie(frameSize); |
145 | void* out = malloc_orDie(dstSize); |
146 | |
147 | size_t inSize = fread_orDie(in, frameSize, fin); |
148 | |
149 | jobs[i].src = in; |
150 | jobs[i].srcSize = inSize; |
151 | jobs[i].dst = out; |
152 | jobs[i].dstSize = dstSize; |
153 | jobs[i].compressionLevel = cLevel; |
154 | jobs[i].done = 0; |
155 | POOL_add(pool, compressFrame, &jobs[i]); |
156 | } |
157 | |
158 | ZSTD_frameLog* fl = ZSTD_seekable_createFrameLog(1); |
159 | if (fl == NULL) { fprintf(stderr, "ZSTD_seekable_createFrameLog() failed \n"); exit(11); } |
160 | for (i = 0; i < numFrames; i++) { |
161 | while (!jobs[i].done) SLEEP(5); /* wake up every 5 milliseconds to check */ |
162 | fwrite_orDie(jobs[i].dst, jobs[i].dstSize, fout); |
163 | free((void*)jobs[i].src); |
164 | free(jobs[i].dst); |
165 | |
166 | size_t ret = ZSTD_seekable_logFrame(fl, jobs[i].dstSize, jobs[i].srcSize, jobs[i].checksum); |
167 | if (ZSTD_isError(ret)) { fprintf(stderr, "ZSTD_seekable_logFrame() error : %s \n", ZSTD_getErrorName(ret)); } |
168 | } |
169 | |
170 | { unsigned char seekTableBuff[1024]; |
171 | ZSTD_outBuffer out = {seekTableBuff, 1024, 0}; |
172 | while (ZSTD_seekable_writeSeekTable(fl, &out) != 0) { |
173 | fwrite_orDie(seekTableBuff, out.pos, fout); |
174 | out.pos = 0; |
175 | } |
176 | fwrite_orDie(seekTableBuff, out.pos, fout); |
177 | } |
178 | |
179 | ZSTD_seekable_freeFrameLog(fl); |
180 | free(jobs); |
181 | fclose_orDie(fout); |
182 | fclose_orDie(fin); |
183 | } |
184 | |
185 | static const char* createOutFilename_orDie(const char* filename) |
186 | { |
187 | size_t const inL = strlen(filename); |
188 | size_t const outL = inL + 5; |
189 | void* outSpace = malloc_orDie(outL); |
190 | memset(outSpace, 0, outL); |
191 | strcat(outSpace, filename); |
192 | strcat(outSpace, ".zst"); |
193 | return (const char*)outSpace; |
194 | } |
195 | |
196 | int main(int argc, const char** argv) { |
197 | const char* const exeName = argv[0]; |
198 | if (argc!=4) { |
199 | printf("wrong arguments\n"); |
200 | printf("usage:\n"); |
201 | printf("%s FILE FRAME_SIZE NB_THREADS\n", exeName); |
202 | return 1; |
203 | } |
204 | |
205 | { const char* const inFileName = argv[1]; |
206 | unsigned const frameSize = (unsigned)atoi(argv[2]); |
207 | int const nbThreads = atoi(argv[3]); |
208 | |
209 | const char* const outFileName = createOutFilename_orDie(inFileName); |
210 | compressFile_orDie(inFileName, outFileName, 5, frameSize, nbThreads); |
211 | } |
212 | |
213 | return 0; |
214 | } |