Line data Source code
1 : /*-------------------------------------------------------------------------
2 : *
3 : * astreamer_zstd.c
4 : *
5 : * Archive streamers that deal with data compressed using zstd.
6 : * astreamer_zstd_compressor applies zstd compression to the input stream,
7 : * and astreamer_zstd_decompressor does the reverse.
8 : *
9 : * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group
10 : *
11 : * IDENTIFICATION
12 : * src/fe_utils/astreamer_zstd.c
13 : *-------------------------------------------------------------------------
14 : */
15 :
16 : #include "postgres_fe.h"
17 :
18 : #include <unistd.h>
19 :
20 : #ifdef USE_ZSTD
21 : #include <zstd.h>
22 : #endif
23 :
24 : #include "common/logging.h"
25 : #include "fe_utils/astreamer.h"
26 :
27 : #ifdef USE_ZSTD
28 :
29 : typedef struct astreamer_zstd_frame
30 : {
31 : astreamer base;
32 :
33 : ZSTD_CCtx *cctx;
34 : ZSTD_DCtx *dctx;
35 : ZSTD_outBuffer zstd_outBuf;
36 : astreamer_decompression_state state;
37 : } astreamer_zstd_frame;
38 :
39 : static void astreamer_zstd_compressor_content(astreamer *streamer,
40 : astreamer_member *member,
41 : const char *data, int len,
42 : astreamer_archive_context context);
43 : static void astreamer_zstd_compressor_finalize(astreamer *streamer);
44 : static void astreamer_zstd_compressor_free(astreamer *streamer);
45 :
46 : static const astreamer_ops astreamer_zstd_compressor_ops = {
47 : .content = astreamer_zstd_compressor_content,
48 : .finalize = astreamer_zstd_compressor_finalize,
49 : .free = astreamer_zstd_compressor_free
50 : };
51 :
52 : static void astreamer_zstd_decompressor_content(astreamer *streamer,
53 : astreamer_member *member,
54 : const char *data, int len,
55 : astreamer_archive_context context);
56 : static void astreamer_zstd_decompressor_finalize(astreamer *streamer);
57 : static void astreamer_zstd_decompressor_free(astreamer *streamer);
58 :
59 : static const astreamer_ops astreamer_zstd_decompressor_ops = {
60 : .content = astreamer_zstd_decompressor_content,
61 : .finalize = astreamer_zstd_decompressor_finalize,
62 : .free = astreamer_zstd_decompressor_free
63 : };
64 : #endif
65 :
66 : /*
67 : * Create a new base backup streamer that performs zstd compression of tar
68 : * blocks.
69 : */
70 : astreamer *
71 0 : astreamer_zstd_compressor_new(astreamer *next, pg_compress_specification *compress)
72 : {
73 : #ifdef USE_ZSTD
74 : astreamer_zstd_frame *streamer;
75 : size_t ret;
76 :
77 : Assert(next != NULL);
78 :
79 : streamer = palloc0_object(astreamer_zstd_frame);
80 :
81 : *((const astreamer_ops **) &streamer->base.bbs_ops) =
82 : &astreamer_zstd_compressor_ops;
83 :
84 : streamer->base.bbs_next = next;
85 : initStringInfo(&streamer->base.bbs_buffer);
86 : enlargeStringInfo(&streamer->base.bbs_buffer, ZSTD_CStreamOutSize());
87 :
88 : streamer->cctx = ZSTD_createCCtx();
89 : if (!streamer->cctx)
90 : pg_fatal("could not create zstd compression context");
91 :
92 : /* Set compression level */
93 : ret = ZSTD_CCtx_setParameter(streamer->cctx, ZSTD_c_compressionLevel,
94 : compress->level);
95 : if (ZSTD_isError(ret))
96 : pg_fatal("could not set zstd compression level to %d: %s",
97 : compress->level, ZSTD_getErrorName(ret));
98 :
99 : /* Set # of workers, if specified */
100 : if ((compress->options & PG_COMPRESSION_OPTION_WORKERS) != 0)
101 : {
102 : /*
103 : * On older versions of libzstd, this option does not exist, and
104 : * trying to set it will fail. Similarly for newer versions if they
105 : * are compiled without threading support.
106 : */
107 : ret = ZSTD_CCtx_setParameter(streamer->cctx, ZSTD_c_nbWorkers,
108 : compress->workers);
109 : if (ZSTD_isError(ret))
110 : pg_fatal("could not set compression worker count to %d: %s",
111 : compress->workers, ZSTD_getErrorName(ret));
112 : }
113 :
114 : if ((compress->options & PG_COMPRESSION_OPTION_LONG_DISTANCE) != 0)
115 : {
116 : ret = ZSTD_CCtx_setParameter(streamer->cctx,
117 : ZSTD_c_enableLongDistanceMatching,
118 : compress->long_distance);
119 : if (ZSTD_isError(ret))
120 : pg_fatal("could not enable long-distance mode: %s",
121 : ZSTD_getErrorName(ret));
122 : }
123 :
124 : /* Initialize the ZSTD output buffer. */
125 : streamer->zstd_outBuf.dst = streamer->base.bbs_buffer.data;
126 : streamer->zstd_outBuf.size = streamer->base.bbs_buffer.maxlen;
127 : streamer->zstd_outBuf.pos = 0;
128 :
129 : return &streamer->base;
130 : #else
131 0 : pg_fatal("this build does not support compression with %s", "ZSTD");
132 : return NULL; /* keep compiler quiet */
133 : #endif
134 : }
135 :
136 : #ifdef USE_ZSTD
137 : /*
138 : * Compress the input data to output buffer.
139 : *
140 : * Find out the compression bound based on input data length for each
141 : * invocation to make sure that output buffer has enough capacity to
142 : * accommodate the compressed data. In case if the output buffer
143 : * capacity falls short of compression bound then forward the content
144 : * of output buffer to next streamer and empty the buffer.
145 : */
146 : static void
147 : astreamer_zstd_compressor_content(astreamer *streamer,
148 : astreamer_member *member,
149 : const char *data, int len,
150 : astreamer_archive_context context)
151 : {
152 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
153 : ZSTD_inBuffer inBuf = {data, len, 0};
154 :
155 : while (inBuf.pos < inBuf.size)
156 : {
157 : size_t yet_to_flush;
158 : size_t max_needed = ZSTD_compressBound(inBuf.size - inBuf.pos);
159 :
160 : /*
161 : * If the output buffer is not left with enough space, send the
162 : * compressed bytes to the next streamer, and empty the buffer.
163 : */
164 : if (mystreamer->zstd_outBuf.size - mystreamer->zstd_outBuf.pos <
165 : max_needed)
166 : {
167 : astreamer_content(mystreamer->base.bbs_next, member,
168 : mystreamer->zstd_outBuf.dst,
169 : mystreamer->zstd_outBuf.pos,
170 : context);
171 :
172 : /* Reset the ZSTD output buffer. */
173 : mystreamer->zstd_outBuf.dst = mystreamer->base.bbs_buffer.data;
174 : mystreamer->zstd_outBuf.size = mystreamer->base.bbs_buffer.maxlen;
175 : mystreamer->zstd_outBuf.pos = 0;
176 : }
177 :
178 : yet_to_flush =
179 : ZSTD_compressStream2(mystreamer->cctx, &mystreamer->zstd_outBuf,
180 : &inBuf, ZSTD_e_continue);
181 :
182 : if (ZSTD_isError(yet_to_flush))
183 : pg_fatal("could not compress data: %s",
184 : ZSTD_getErrorName(yet_to_flush));
185 : }
186 : }
187 :
188 : /*
189 : * End-of-stream processing.
190 : */
191 : static void
192 : astreamer_zstd_compressor_finalize(astreamer *streamer)
193 : {
194 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
195 : size_t yet_to_flush;
196 :
197 : do
198 : {
199 : ZSTD_inBuffer in = {NULL, 0, 0};
200 : size_t max_needed = ZSTD_compressBound(0);
201 :
202 : /*
203 : * If the output buffer is not left with enough space, send the
204 : * compressed bytes to the next streamer, and empty the buffer.
205 : */
206 : if (mystreamer->zstd_outBuf.size - mystreamer->zstd_outBuf.pos <
207 : max_needed)
208 : {
209 : astreamer_content(mystreamer->base.bbs_next, NULL,
210 : mystreamer->zstd_outBuf.dst,
211 : mystreamer->zstd_outBuf.pos,
212 : ASTREAMER_UNKNOWN);
213 :
214 : /* Reset the ZSTD output buffer. */
215 : mystreamer->zstd_outBuf.dst = mystreamer->base.bbs_buffer.data;
216 : mystreamer->zstd_outBuf.size = mystreamer->base.bbs_buffer.maxlen;
217 : mystreamer->zstd_outBuf.pos = 0;
218 : }
219 :
220 : yet_to_flush = ZSTD_compressStream2(mystreamer->cctx,
221 : &mystreamer->zstd_outBuf,
222 : &in, ZSTD_e_end);
223 :
224 : if (ZSTD_isError(yet_to_flush))
225 : pg_fatal("could not compress data: %s",
226 : ZSTD_getErrorName(yet_to_flush));
227 :
228 : } while (yet_to_flush > 0);
229 :
230 : /* Make sure to pass any remaining bytes to the next streamer. */
231 : if (mystreamer->zstd_outBuf.pos > 0)
232 : astreamer_content(mystreamer->base.bbs_next, NULL,
233 : mystreamer->zstd_outBuf.dst,
234 : mystreamer->zstd_outBuf.pos,
235 : ASTREAMER_UNKNOWN);
236 :
237 : astreamer_finalize(mystreamer->base.bbs_next);
238 : }
239 :
240 : /*
241 : * Free memory.
242 : */
243 : static void
244 : astreamer_zstd_compressor_free(astreamer *streamer)
245 : {
246 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
247 :
248 : astreamer_free(streamer->bbs_next);
249 : ZSTD_freeCCtx(mystreamer->cctx);
250 : pfree(streamer->bbs_buffer.data);
251 : pfree(streamer);
252 : }
253 : #endif
254 :
255 : /*
256 : * Create a new base backup streamer that performs decompression of zstd
257 : * compressed blocks.
258 : */
259 : astreamer *
260 0 : astreamer_zstd_decompressor_new(astreamer *next)
261 : {
262 : #ifdef USE_ZSTD
263 : astreamer_zstd_frame *streamer;
264 :
265 : Assert(next != NULL);
266 :
267 : streamer = palloc0_object(astreamer_zstd_frame);
268 : *((const astreamer_ops **) &streamer->base.bbs_ops) =
269 : &astreamer_zstd_decompressor_ops;
270 :
271 : streamer->base.bbs_next = next;
272 : initStringInfo(&streamer->base.bbs_buffer);
273 : enlargeStringInfo(&streamer->base.bbs_buffer, ZSTD_DStreamOutSize());
274 :
275 : streamer->dctx = ZSTD_createDCtx();
276 : if (!streamer->dctx)
277 : pg_fatal("could not create zstd decompression context");
278 :
279 : /* Initialize the ZSTD output buffer. */
280 : streamer->zstd_outBuf.dst = streamer->base.bbs_buffer.data;
281 : streamer->zstd_outBuf.size = streamer->base.bbs_buffer.maxlen;
282 : streamer->zstd_outBuf.pos = 0;
283 :
284 : streamer->state = ASTREAMER_STREAM_NEW;
285 : return &streamer->base;
286 : #else
287 0 : pg_fatal("this build does not support compression with %s", "ZSTD");
288 : return NULL; /* keep compiler quiet */
289 : #endif
290 : }
291 :
292 : #ifdef USE_ZSTD
293 : /*
294 : * Decompress the input data to output buffer until we run out of input
295 : * data. Each time the output buffer is full, pass on the decompressed data
296 : * to the next streamer.
297 : */
298 : static void
299 : astreamer_zstd_decompressor_content(astreamer *streamer,
300 : astreamer_member *member,
301 : const char *data, int len,
302 : astreamer_archive_context context)
303 : {
304 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
305 : ZSTD_inBuffer inBuf = {data, len, 0};
306 :
307 : while (inBuf.pos < inBuf.size)
308 : {
309 : size_t ret;
310 :
311 : /*
312 : * If output buffer is full then forward the content to next streamer
313 : * and update the output buffer.
314 : */
315 : if (mystreamer->zstd_outBuf.pos >= mystreamer->zstd_outBuf.size)
316 : {
317 : astreamer_content(mystreamer->base.bbs_next, member,
318 : mystreamer->zstd_outBuf.dst,
319 : mystreamer->zstd_outBuf.pos,
320 : context);
321 :
322 : /* Reset the ZSTD output buffer. */
323 : mystreamer->zstd_outBuf.dst = mystreamer->base.bbs_buffer.data;
324 : mystreamer->zstd_outBuf.size = mystreamer->base.bbs_buffer.maxlen;
325 : mystreamer->zstd_outBuf.pos = 0;
326 : }
327 :
328 : ret = ZSTD_decompressStream(mystreamer->dctx,
329 : &mystreamer->zstd_outBuf, &inBuf);
330 :
331 : if (ZSTD_isError(ret))
332 : pg_fatal("could not decompress data: %s",
333 : ZSTD_getErrorName(ret));
334 :
335 : /* The frame is only done when ZSTD_decompressStream returns 0 */
336 : if (ret)
337 : mystreamer->state = ASTREAMER_FRAME_INCOMPLETE;
338 : else
339 : mystreamer->state = ASTREAMER_FRAME_COMPLETE;
340 : }
341 : }
342 :
343 : /*
344 : * End-of-stream processing.
345 : */
346 : static void
347 : astreamer_zstd_decompressor_finalize(astreamer *streamer)
348 : {
349 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
350 :
351 : /*
352 : * A full output buffer with a positive return value might leave data in
353 : * zstd's internal buffers. Call the decompressor with empty input until
354 : * it has flushed that data.
355 : */
356 : while (mystreamer->state == ASTREAMER_FRAME_INCOMPLETE &&
357 : mystreamer->zstd_outBuf.pos == mystreamer->zstd_outBuf.size)
358 : {
359 : ZSTD_inBuffer empty = {NULL, 0, 0};
360 : size_t ret;
361 :
362 : astreamer_content(mystreamer->base.bbs_next, NULL,
363 : mystreamer->zstd_outBuf.dst,
364 : mystreamer->zstd_outBuf.pos, ASTREAMER_UNKNOWN);
365 :
366 : mystreamer->zstd_outBuf.dst = mystreamer->base.bbs_buffer.data;
367 : mystreamer->zstd_outBuf.size = mystreamer->base.bbs_buffer.maxlen;
368 : mystreamer->zstd_outBuf.pos = 0;
369 :
370 : ret = ZSTD_decompressStream(mystreamer->dctx,
371 : &mystreamer->zstd_outBuf, &empty);
372 :
373 : if (ZSTD_isError(ret))
374 : pg_fatal("could not decompress data: %s",
375 : ZSTD_getErrorName(ret));
376 :
377 : /* The frame is only done when ZSTD_decompressStream returns 0 */
378 : if (ret)
379 : mystreamer->state = ASTREAMER_FRAME_INCOMPLETE;
380 : else
381 : mystreamer->state = ASTREAMER_FRAME_COMPLETE;
382 : }
383 :
384 : if (unlikely(mystreamer->state == ASTREAMER_STREAM_NEW))
385 : pg_fatal("could not decompress data: compressed stream is empty");
386 : else if (mystreamer->state != ASTREAMER_FRAME_COMPLETE)
387 : pg_fatal("could not decompress data: compressed stream is incomplete");
388 :
389 : /*
390 : * End of the stream, if there is some pending data in output buffers then
391 : * we must forward it to next streamer.
392 : */
393 : if (mystreamer->zstd_outBuf.pos > 0)
394 : astreamer_content(mystreamer->base.bbs_next, NULL,
395 : mystreamer->base.bbs_buffer.data,
396 : mystreamer->zstd_outBuf.pos,
397 : ASTREAMER_UNKNOWN);
398 :
399 : astreamer_finalize(mystreamer->base.bbs_next);
400 : }
401 :
402 : /*
403 : * Free memory.
404 : */
405 : static void
406 : astreamer_zstd_decompressor_free(astreamer *streamer)
407 : {
408 : astreamer_zstd_frame *mystreamer = (astreamer_zstd_frame *) streamer;
409 :
410 : astreamer_free(streamer->bbs_next);
411 : ZSTD_freeDCtx(mystreamer->dctx);
412 : pfree(streamer->bbs_buffer.data);
413 : pfree(streamer);
414 : }
415 : #endif
|
