1 // SPDX-License-Identifier: CDDL-1.0
2 /*
3 * This file and its contents are supplied under the terms of the
4 * Common Development and Distribution License ("CDDL"), version 1.0.
5 * You may only use this file in accordance with the terms of version
6 * 1.0 of the CDDL.
7 *
8 * A full copy of the text of the CDDL should have accompanied this
9 * source. A copy of the CDDL is also available via the Internet at
10 * https://opensource.org/license/CDDL-1.0.
11 */
12
13 /*
14 * Copyright 2022 Axcient. All rights reserved.
15 * Use is subject to license terms.
16 *
17 * Copyright (c) 2022 by Delphix. All rights reserved.
18 * Copyright (c) 2024, Klara, Inc.
19 * Copyright (c) 2026 by Garth Snyder
20 */
21
22 #include <assert.h>
23 #include <err.h>
24 #include <stdint.h>
25 #include <stdio.h>
26 #include <stdlib.h>
27 #include <string.h>
28 #include <sys/dmu.h>
29 #include <sys/zfs_ioctl.h>
30 #include <sys/zio_compress.h>
31 #include <sys/zstd/zstd.h>
32 #include <unistd.h>
33 #include <sys/stdtypes.h>
34
35 #include "zstream.h"
36 #include "zstream_chain.h"
37 #include "zstream_modules.h"
38 #include "zstream_queue.h"
39 #include "zstream_recompress.h"
40
41 #define MAX_COMPRESSION_STEPS 4
42
43 static compression_spec_t specs[MAX_COMPRESSION_STEPS];
44 static int next_spec = 0;
45
46 /*
47 * Item is known to be a DRR_WRITE packet. Determine whether current
48 * compression is compatible with desired compression and whether the
49 * current record is modifiable at all.
50 */
51 static boolean_t
needs_modification(drr_packet_t * item,compression_spec_t * target)52 needs_modification(drr_packet_t *item, compression_spec_t *target)
53 {
54 dmu_replay_record_t *drr = &item->dp_drr;
55 struct drr_write *drrw = &drr->drr_u.drr_write;
56 enum zio_compress ctype = drrw->drr_compressiontype;
57 uint8_t cur_level;
58
59 /*
60 * Do not modify metadata records. It's a general stream invariant
61 * that metadata is never compressed. See comments at
62 * dmu_receive.c:flush_write_batch_impl().
63 */
64 if (DMU_OT_IS_METADATA(drrw->drr_type)) {
65 return (B_FALSE);
66 }
67 boolean_t ctype_uncompressed = ctype_is_uncompressed(ctype);
68 if (target == NULL) {
69 return (!ctype_uncompressed && !write_is_encrypted(drrw));
70 }
71 boolean_t target_uncompressed = ctype_is_uncompressed(target->cs_type);
72 if (target_uncompressed && ctype_uncompressed) {
73 return (B_FALSE);
74 }
75 /*
76 * In order to recompress an encrypted block, you have to decrypt,
77 * decompress, recompress, and re-encrypt. That can be a future
78 * enhancement (along with decryption or re-encryption), but for now
79 * we skip encrypted blocks.
80 */
81 if (write_is_encrypted(drrw)) {
82 return (B_FALSE);
83 }
84 if (ctype != target->cs_type) {
85 return (B_TRUE);
86 }
87 if (target->cs_type == ZIO_COMPRESS_ZSTD) {
88 cur_level = zfs_get_hdrlevel((void *)item->dp_payload);
89 if (target->cs_level == ZIO_COMPLEVEL_DEFAULT) {
90 return (cur_level != ZIO_ZSTD_LEVEL_DEFAULT);
91 }
92 return (target->cs_level != cur_level);
93 }
94 return (B_FALSE);
95 }
96
97 static boolean_t
needs_compression(drr_packet_t * item,compression_spec_t * context)98 needs_compression(drr_packet_t *item, compression_spec_t *context)
99 {
100 return (needs_modification(item, context));
101 }
102
103 /*
104 * Don't decompress packets that aren't compressed. And don't decompress
105 * them if their ultimate fate is to be recompressed using the compression
106 * profile that's already in use.
107 */
108 static boolean_t
needs_decompression(drr_packet_t * item,compression_spec_t * context)109 needs_decompression(drr_packet_t *item, compression_spec_t *context)
110 {
111 dmu_replay_record_t *drr = &item->dp_drr;
112 struct drr_write *drrw = &drr->drr_u.drr_write;
113 enum zio_compress ctype = drrw->drr_compressiontype;
114
115 if (ctype_is_uncompressed(ctype))
116 return (B_FALSE);
117 return (needs_modification(item, context));
118 }
119
120 /*
121 * We can ignore the context here because it's already been evaluated by the
122 * cost function. If the cost function returned something other than zero,
123 * we have to decompress.
124 */
125 static void
chain_decompress_writes(queue_item_t * item_in,void * context)126 chain_decompress_writes(queue_item_t *item_in, void *context)
127 {
128 (void) context;
129 drr_packet_t *item = (drr_packet_t *)item_in;
130
131 dmu_replay_record_t *drr = &item->dp_drr;
132 struct drr_write *drrw = &drr->drr_u.drr_write;
133 uint8_t *debuff;
134
135 VERIFY3U(drr->drr_type, ==, DRR_WRITE);
136 debuff = decompress_buffer(item->dp_payload, item->dp_payload_size,
137 drrw->drr_logical_size, drrw->drr_compressiontype);
138 if (debuff == NULL) {
139 errx(4, "decompression type %d failed for ino %llu offset %llu",
140 drrw->drr_compressiontype,
141 (u_longlong_t)drrw->drr_object,
142 (u_longlong_t)drrw->drr_offset);
143 }
144 free(item->dp_payload);
145 item->dp_payload = debuff;
146 item->dp_payload_size = drrw->drr_logical_size;
147 drrw->drr_compressed_size = 0;
148 drrw->drr_compressiontype = 0;
149 }
150
151 /*
152 * As with chain_decompress_writes(), all the important decisions were made
153 * by the cost function. If we're here, we need to compress.
154 */
155 static void
chain_compress_writes(queue_item_t * item_in,void * context_in)156 chain_compress_writes(queue_item_t *item_in, void *context_in)
157 {
158 drr_packet_t *item = (drr_packet_t *)item_in;
159 compression_spec_t *context = (compression_spec_t *)context_in;
160
161 dmu_replay_record_t *drr = &item->dp_drr;
162
163 struct drr_write *drrw = &drr->drr_u.drr_write;
164 enum zio_compress ctype = drrw->drr_compressiontype;
165 uint8_t *cbuff;
166 size_t csize;
167
168 VERIFY3U(drr->drr_type, ==, DRR_WRITE);
169 VERIFY3B(ctype_is_uncompressed(ctype), ==, B_TRUE);
170 cbuff = compress_buffer(item->dp_payload, item->dp_payload_size,
171 *context, &csize);
172
173 if (cbuff == NULL) {
174 drrw->drr_compressiontype = 0;
175 drrw->drr_compressed_size = 0;
176 } else {
177 free(item->dp_payload);
178 item->dp_payload = cbuff;
179 item->dp_payload_size = csize;
180 drrw->drr_compressed_size = csize;
181 drrw->drr_compressiontype = context->cs_type;
182 }
183 }
184
185 /*
186 * A cost of zero waives processing for the current item. If we want to
187 * process it, the cost will always be item->dp_payload_size. So in these
188 * two cost functions, we're mostly determining which packets need
189 * attention. A packet that's already compressed with the target compression
190 * profile can be ignored.
191 */
192 static size_t
chain_compress_cost(queue_item_t * item_in,void * context_in)193 chain_compress_cost(queue_item_t *item_in, void *context_in)
194 {
195 compression_spec_t *context = (compression_spec_t *)context_in;
196 drr_packet_t *item = (drr_packet_t *)item_in;
197 dmu_replay_record_t *drr = &item->dp_drr;
198
199 if (drr->drr_type != DRR_WRITE) {
200 return (0);
201 }
202 struct drr_write *drrw = &drr->drr_u.drr_write;
203 return (needs_compression(item, context) ? drrw->drr_logical_size : 0);
204 }
205
206 /*
207 * Don't decompress packets that aren't compressed. And don't decompress
208 * them if their ultimate fate is to be recompressed using the compression
209 * profile that's already in use.
210 */
211 static size_t
chain_decompress_cost(queue_item_t * item_in,void * context_in)212 chain_decompress_cost(queue_item_t *item_in, void *context_in)
213 {
214 compression_spec_t *context = (compression_spec_t *)context_in;
215 drr_packet_t *item = (drr_packet_t *)item_in;
216 dmu_replay_record_t *drr = &item->dp_drr;
217
218 if (drr->drr_type != DRR_WRITE)
219 return (0);
220
221 struct drr_write *drrw = &drr->drr_u.drr_write;
222 enum zio_compress ctype = drrw->drr_compressiontype;
223
224 if (ctype_is_uncompressed(ctype))
225 return (0);
226
227 return (needs_decompression(item, context) ? item->dp_payload_size : 0);
228 }
229
230 /*
231 * Decompress writes, but only if they don't match a target compression
232 * type. Pass NULL to uncompress unconditionally (if not already
233 * uncompressed).
234 */
235 chain_step_t
parallel_decompress_writes(compression_spec_t * target)236 parallel_decompress_writes(compression_spec_t *target)
237 {
238 int this_spec = next_spec++ % MAX_COMPRESSION_STEPS;
239 compression_spec_t *context = &specs[this_spec];
240
241 if (target == NULL) {
242 context = NULL;
243 } else {
244 *context = *target;
245 }
246 chain_step_t step = {
247 .cs_type = CS_PARALLEL,
248 .cs_in_size = sizeof (drr_packet_t),
249 .cs_out_size = sizeof (drr_packet_t),
250 .cs_context = context,
251 .cs_parallel = {
252 .queue_length = 256,
253 .batch_budget = 256 * 1024,
254 .process = chain_decompress_writes,
255 .cost = chain_decompress_cost
256 }
257 };
258 return (step);
259 }
260
261 chain_step_t
parallel_compress_writes(compression_spec_t * target)262 parallel_compress_writes(compression_spec_t *target)
263 {
264 int this_spec = next_spec++ % MAX_COMPRESSION_STEPS;
265 compression_spec_t *context = &specs[this_spec];
266
267 VERIFY3P(target, !=, NULL);
268 *context = *target;
269
270 chain_step_t step = {
271 .cs_type = CS_PARALLEL,
272 .cs_in_size = sizeof (drr_packet_t),
273 .cs_out_size = sizeof (drr_packet_t),
274 .cs_context = context,
275 .cs_parallel = {
276 .queue_length = 1024,
277 .batch_budget = 32 * 1024,
278 .process = chain_compress_writes,
279 .cost = chain_compress_cost
280 }
281 };
282 return (step);
283 }
284
285 /*
286 * Keep DRR_BEGIN feature flags consistent with the WRITE payloads we emit.
287 * Compressed WRITEs require DMU_BACKUP_FEATURE_COMPRESSED (and LZ4/ZSTD as
288 * appropriate). Recompressing to off clears COMPRESSED for non-raw streams;
289 * raw streams may retain encrypted compressed WRITEs that recompress skips.
290 * LZ4/ZSTD are cleared only when neither EMBED_DATA nor RAW remains, since
291 * recompress does not rewrite those record types.
292 */
293 static disposition_t
chain_update_compress_features(void * item_in,void * context_in)294 chain_update_compress_features(void *item_in, void *context_in)
295 {
296 drr_packet_t *item = (drr_packet_t *)item_in;
297 compression_spec_t *spec = (compression_spec_t *)context_in;
298 struct drr_begin *drrb;
299 uint64_t flags;
300
301 if (item == NULL)
302 return (D_OK);
303
304 if (item->dp_drr.drr_type != DRR_BEGIN)
305 return (D_OK);
306
307 drrb = &item->dp_drr.drr_u.drr_begin;
308 flags = DMU_GET_FEATUREFLAGS(drrb->drr_versioninfo);
309
310 if (ctype_is_uncompressed(spec->cs_type)) {
311 if (!(flags & DMU_BACKUP_FEATURE_RAW))
312 flags &= ~DMU_BACKUP_FEATURE_COMPRESSED;
313 if ((flags & (DMU_BACKUP_FEATURE_EMBED_DATA |
314 DMU_BACKUP_FEATURE_RAW)) == 0) {
315 flags &= ~(DMU_BACKUP_FEATURE_LZ4 |
316 DMU_BACKUP_FEATURE_ZSTD);
317 }
318 } else {
319 flags |= DMU_BACKUP_FEATURE_COMPRESSED;
320 if (spec->cs_type == ZIO_COMPRESS_ZSTD) {
321 flags |= DMU_BACKUP_FEATURE_ZSTD;
322 } else if (spec->cs_type >= ZIO_COMPRESS_LZ4) {
323 flags |= DMU_BACKUP_FEATURE_LZ4;
324 }
325 }
326
327 DMU_SET_FEATUREFLAGS(drrb->drr_versioninfo, flags);
328 return (D_OK);
329 }
330
331 static chain_step_t
serial_update_compress_features(compression_spec_t * target)332 serial_update_compress_features(compression_spec_t *target)
333 {
334 int this_spec = next_spec++ % MAX_COMPRESSION_STEPS;
335 compression_spec_t *context = &specs[this_spec];
336
337 VERIFY3P(target, !=, NULL);
338 *context = *target;
339
340 chain_step_t step = {
341 .cs_type = CS_SERIAL,
342 .cs_in_size = sizeof (drr_packet_t),
343 .cs_out_size = sizeof (drr_packet_t),
344 .cs_context = context,
345 .cs_serial = {
346 .process = chain_update_compress_features,
347 }
348 };
349 return (step);
350 }
351
352 int
zstream_do_recompress(int argc,char * argv[])353 zstream_do_recompress(int argc, char *argv[])
354 {
355 int c;
356 int level = ZIO_COMPLEVEL_DEFAULT;
357 int num_threads = 0;
358
359 chain_attrs_t attrs = { .ca_command_opts = CA_FORBID_DEDUP };
360
361 while ((c = getopt(argc, argv, "t:l:")) != -1) {
362 switch (c) {
363 case 'l':
364 if (sscanf(optarg, "%d", &level) != 1) {
365 warnx("failed to parse level '%s'", optarg);
366 zstream_usage();
367 }
368 break;
369 case 't':
370 if (sscanf(optarg, "%d", &num_threads) != 1) {
371 warnx("failed to parse num_threads '%s'",
372 optarg);
373 zstream_usage();
374 }
375 zstream_queue_set_num_threads(num_threads);
376 break;
377 case '?':
378 warnx("invalid option '%c'", optopt);
379 zstream_usage();
380 }
381 }
382
383 argc -= optind;
384 argv += optind;
385
386 if (argc != 1)
387 zstream_usage();
388
389 compression_spec_t spec = { .cs_level = level };
390 if (strcmp(argv[0], "off") == 0) {
391 spec.cs_type = ZIO_COMPRESS_OFF;
392 } else {
393 enum zio_compress ct;
394 for (ct = 0; ct < ZIO_COMPRESS_FUNCTIONS; ct++) {
395 const char *ci_name = zio_compress_table[ct].ci_name;
396 if (strcmp(argv[0], ci_name) == 0)
397 break;
398 }
399 if (ct == ZIO_COMPRESS_FUNCTIONS || ctype_is_uncompressed(ct)) {
400 errx(2, "invalid compression type %s", argv[0]);
401 }
402 spec.cs_type = ct;
403 }
404
405 zstream_chain_t recompress_chain = {
406 STANDARD_INPUT_STACK(NULL),
407 parallel_decompress_writes(&spec),
408 parallel_compress_writes(&spec),
409 serial_update_compress_features(&spec),
410 STANDARD_OUTPUT_STACK(NULL)
411 };
412
413 zstream_chain_exec(recompress_chain, &attrs);
414 return (0);
415 }
416