xref: /freebsd/sys/contrib/openzfs/module/zstd/zfs_zstd.c (revision 2f10ffc003be396f3fc23cd2888023896560252b)
1 // SPDX-License-Identifier: BSD-3-Clause
2 /*
3  * BSD 3-Clause New License (https://spdx.org/licenses/BSD-3-Clause.html)
4  *
5  * Redistribution and use in source and binary forms, with or without
6  * modification, are permitted provided that the following conditions are met:
7  *
8  * 1. Redistributions of source code must retain the above copyright notice,
9  * this list of conditions and the following disclaimer.
10  *
11  * 2. Redistributions in binary form must reproduce the above copyright notice,
12  * this list of conditions and the following disclaimer in the documentation
13  * and/or other materials provided with the distribution.
14  *
15  * 3. Neither the name of the copyright holder nor the names of its
16  * contributors may be used to endorse or promote products derived from this
17  * software without specific prior written permission.
18  *
19  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29  * POSSIBILITY OF SUCH DAMAGE.
30  */
31 
32 /*
33  * Copyright (c) 2016-2018, Klara Inc.
34  * Copyright (c) 2016-2018, Allan Jude
35  * Copyright (c) 2018-2020, Sebastian Gottschall
36  * Copyright (c) 2019-2020, Michael Niewöhner
37  * Copyright (c) 2020, The FreeBSD Foundation [1]
38  *
39  * [1] Portions of this software were developed by Allan Jude
40  *     under sponsorship from the FreeBSD Foundation.
41  */
42 
43 #include <sys/param.h>
44 #include <sys/sysmacros.h>
45 #include <sys/zfs_context.h>
46 #include <sys/zio_compress.h>
47 #include <sys/spa.h>
48 #include <sys/zstd/zstd.h>
49 
50 #define	ZSTD_STATIC_LINKING_ONLY
51 #include "lib/zstd.h"
52 #include "lib/zstd_errors.h"
53 
54 #ifndef IN_LIBSA
55 static uint_t zstd_earlyabort_pass = 1;
56 static int zstd_cutoff_level = ZIO_ZSTD_LEVEL_3;
57 static unsigned int zstd_abort_size = (128 * 1024);
58 #endif
59 
60 #ifdef IN_BASE
61 int zfs_zstd_decompress_buf(void *, void *, size_t, size_t, int);
62 #endif
63 
64 static kstat_t *zstd_ksp = NULL;
65 
66 typedef struct zstd_stats {
67 	kstat_named_t	zstd_stat_alloc_fail;
68 	kstat_named_t	zstd_stat_alloc_fallback;
69 	kstat_named_t	zstd_stat_com_alloc_fail;
70 	kstat_named_t	zstd_stat_dec_alloc_fail;
71 	kstat_named_t	zstd_stat_com_inval;
72 	kstat_named_t	zstd_stat_dec_inval;
73 	kstat_named_t	zstd_stat_dec_header_inval;
74 	kstat_named_t	zstd_stat_com_fail;
75 	kstat_named_t	zstd_stat_dec_fail;
76 	/*
77 	 * LZ4 first-pass early abort verdict
78 	 */
79 	kstat_named_t	zstd_stat_lz4pass_allowed;
80 	kstat_named_t	zstd_stat_lz4pass_rejected;
81 	/*
82 	 * zstd-1 second-pass early abort verdict
83 	 */
84 	kstat_named_t	zstd_stat_zstdpass_allowed;
85 	kstat_named_t	zstd_stat_zstdpass_rejected;
86 	/*
87 	 * We excluded this from early abort for some reason
88 	 */
89 	kstat_named_t	zstd_stat_passignored;
90 	kstat_named_t	zstd_stat_passignored_size;
91 	kstat_named_t	zstd_stat_buffers;
92 	kstat_named_t	zstd_stat_size;
93 } zstd_stats_t;
94 
95 static zstd_stats_t zstd_stats = {
96 	{ "alloc_fail",			KSTAT_DATA_UINT64 },
97 	{ "alloc_fallback",		KSTAT_DATA_UINT64 },
98 	{ "compress_alloc_fail",	KSTAT_DATA_UINT64 },
99 	{ "decompress_alloc_fail",	KSTAT_DATA_UINT64 },
100 	{ "compress_level_invalid",	KSTAT_DATA_UINT64 },
101 	{ "decompress_level_invalid",	KSTAT_DATA_UINT64 },
102 	{ "decompress_header_invalid",	KSTAT_DATA_UINT64 },
103 	{ "compress_failed",		KSTAT_DATA_UINT64 },
104 	{ "decompress_failed",		KSTAT_DATA_UINT64 },
105 	{ "lz4pass_allowed",		KSTAT_DATA_UINT64 },
106 	{ "lz4pass_rejected",		KSTAT_DATA_UINT64 },
107 	{ "zstdpass_allowed",		KSTAT_DATA_UINT64 },
108 	{ "zstdpass_rejected",		KSTAT_DATA_UINT64 },
109 	{ "passignored",		KSTAT_DATA_UINT64 },
110 	{ "passignored_size",		KSTAT_DATA_UINT64 },
111 	{ "buffers",			KSTAT_DATA_UINT64 },
112 	{ "size",			KSTAT_DATA_UINT64 },
113 };
114 
115 #ifdef _KERNEL
116 static int
117 kstat_zstd_update(kstat_t *ksp, int rw)
118 {
119 	ASSERT(ksp != NULL);
120 
121 	if (rw == KSTAT_WRITE && ksp == zstd_ksp) {
122 		ZSTDSTAT_ZERO(zstd_stat_alloc_fail);
123 		ZSTDSTAT_ZERO(zstd_stat_alloc_fallback);
124 		ZSTDSTAT_ZERO(zstd_stat_com_alloc_fail);
125 		ZSTDSTAT_ZERO(zstd_stat_dec_alloc_fail);
126 		ZSTDSTAT_ZERO(zstd_stat_com_inval);
127 		ZSTDSTAT_ZERO(zstd_stat_dec_inval);
128 		ZSTDSTAT_ZERO(zstd_stat_dec_header_inval);
129 		ZSTDSTAT_ZERO(zstd_stat_com_fail);
130 		ZSTDSTAT_ZERO(zstd_stat_dec_fail);
131 		ZSTDSTAT_ZERO(zstd_stat_lz4pass_allowed);
132 		ZSTDSTAT_ZERO(zstd_stat_lz4pass_rejected);
133 		ZSTDSTAT_ZERO(zstd_stat_zstdpass_allowed);
134 		ZSTDSTAT_ZERO(zstd_stat_zstdpass_rejected);
135 		ZSTDSTAT_ZERO(zstd_stat_passignored);
136 		ZSTDSTAT_ZERO(zstd_stat_passignored_size);
137 	}
138 
139 	return (0);
140 }
141 #endif
142 
143 /* Enums describing the allocator type specified by kmem_type in zstd_kmem */
144 enum zstd_kmem_type {
145 	ZSTD_KMEM_UNKNOWN = 0,
146 	/* Allocation type using kmem_vmalloc */
147 	ZSTD_KMEM_DEFAULT,
148 	/* Pool based allocation using mempool_alloc */
149 	ZSTD_KMEM_POOL,
150 	/* Reserved fallback memory for decompression only */
151 	ZSTD_KMEM_DCTX,
152 	ZSTD_KMEM_COUNT,
153 };
154 
155 /* Structure for pooled memory objects */
156 struct zstd_pool {
157 	void *mem;
158 	size_t size;
159 	kmutex_t barrier;
160 	hrtime_t timeout;
161 };
162 
163 /* Global structure for handling memory allocations */
164 struct zstd_kmem {
165 	enum zstd_kmem_type kmem_type;
166 	size_t kmem_size;
167 	struct zstd_pool *pool;
168 };
169 
170 /* Fallback memory structure used for decompression only if memory runs out */
171 struct zstd_fallback_mem {
172 	size_t mem_size;
173 	void *mem;
174 	kmutex_t barrier;
175 };
176 
177 struct zstd_levelmap {
178 	int16_t zstd_level;
179 	enum zio_zstd_levels level;
180 };
181 
182 /*
183  * ZSTD memory handlers
184  *
185  * For decompression we use a different handler which also provides fallback
186  * memory allocation in case memory runs out.
187  *
188  * The ZSTD handlers were split up for the most simplified implementation.
189  */
190 #ifndef IN_LIBSA
191 static void *zstd_alloc(void *opaque, size_t size);
192 #endif
193 static void *zstd_dctx_alloc(void *opaque, size_t size);
194 static void zstd_free(void *opaque, void *ptr);
195 
196 #ifndef IN_LIBSA
197 /* Compression memory handler */
198 static const ZSTD_customMem zstd_malloc = {
199 	zstd_alloc,
200 	zstd_free,
201 	NULL,
202 };
203 #endif
204 
205 /* Decompression memory handler */
206 static const ZSTD_customMem zstd_dctx_malloc = {
207 	zstd_dctx_alloc,
208 	zstd_free,
209 	NULL,
210 };
211 
212 /* Level map for converting ZFS internal levels to ZSTD levels and vice versa */
213 static struct zstd_levelmap zstd_levels[] = {
214 	{ZIO_ZSTD_LEVEL_1, ZIO_ZSTD_LEVEL_1},
215 	{ZIO_ZSTD_LEVEL_2, ZIO_ZSTD_LEVEL_2},
216 	{ZIO_ZSTD_LEVEL_3, ZIO_ZSTD_LEVEL_3},
217 	{ZIO_ZSTD_LEVEL_4, ZIO_ZSTD_LEVEL_4},
218 	{ZIO_ZSTD_LEVEL_5, ZIO_ZSTD_LEVEL_5},
219 	{ZIO_ZSTD_LEVEL_6, ZIO_ZSTD_LEVEL_6},
220 	{ZIO_ZSTD_LEVEL_7, ZIO_ZSTD_LEVEL_7},
221 	{ZIO_ZSTD_LEVEL_8, ZIO_ZSTD_LEVEL_8},
222 	{ZIO_ZSTD_LEVEL_9, ZIO_ZSTD_LEVEL_9},
223 	{ZIO_ZSTD_LEVEL_10, ZIO_ZSTD_LEVEL_10},
224 	{ZIO_ZSTD_LEVEL_11, ZIO_ZSTD_LEVEL_11},
225 	{ZIO_ZSTD_LEVEL_12, ZIO_ZSTD_LEVEL_12},
226 	{ZIO_ZSTD_LEVEL_13, ZIO_ZSTD_LEVEL_13},
227 	{ZIO_ZSTD_LEVEL_14, ZIO_ZSTD_LEVEL_14},
228 	{ZIO_ZSTD_LEVEL_15, ZIO_ZSTD_LEVEL_15},
229 	{ZIO_ZSTD_LEVEL_16, ZIO_ZSTD_LEVEL_16},
230 	{ZIO_ZSTD_LEVEL_17, ZIO_ZSTD_LEVEL_17},
231 	{ZIO_ZSTD_LEVEL_18, ZIO_ZSTD_LEVEL_18},
232 	{ZIO_ZSTD_LEVEL_19, ZIO_ZSTD_LEVEL_19},
233 	{-1, ZIO_ZSTD_LEVEL_FAST_1},
234 	{-2, ZIO_ZSTD_LEVEL_FAST_2},
235 	{-3, ZIO_ZSTD_LEVEL_FAST_3},
236 	{-4, ZIO_ZSTD_LEVEL_FAST_4},
237 	{-5, ZIO_ZSTD_LEVEL_FAST_5},
238 	{-6, ZIO_ZSTD_LEVEL_FAST_6},
239 	{-7, ZIO_ZSTD_LEVEL_FAST_7},
240 	{-8, ZIO_ZSTD_LEVEL_FAST_8},
241 	{-9, ZIO_ZSTD_LEVEL_FAST_9},
242 	{-10, ZIO_ZSTD_LEVEL_FAST_10},
243 	{-20, ZIO_ZSTD_LEVEL_FAST_20},
244 	{-30, ZIO_ZSTD_LEVEL_FAST_30},
245 	{-40, ZIO_ZSTD_LEVEL_FAST_40},
246 	{-50, ZIO_ZSTD_LEVEL_FAST_50},
247 	{-60, ZIO_ZSTD_LEVEL_FAST_60},
248 	{-70, ZIO_ZSTD_LEVEL_FAST_70},
249 	{-80, ZIO_ZSTD_LEVEL_FAST_80},
250 	{-90, ZIO_ZSTD_LEVEL_FAST_90},
251 	{-100, ZIO_ZSTD_LEVEL_FAST_100},
252 	{-500, ZIO_ZSTD_LEVEL_FAST_500},
253 	{-1000, ZIO_ZSTD_LEVEL_FAST_1000},
254 };
255 
256 /*
257  * This variable represents the maximum count of the pool based on the number
258  * of CPUs plus some buffer. We default to cpu count * 4, see init_zstd.
259  */
260 static int pool_count = 16;
261 
262 #define	ZSTD_POOL_MAX		pool_count
263 #define	ZSTD_POOL_TIMEOUT	60 * 2
264 
265 static struct zstd_fallback_mem zstd_dctx_fallback;
266 static struct zstd_pool *zstd_mempool_cctx;
267 static struct zstd_pool *zstd_mempool_dctx;
268 
269 /*
270  * The library zstd code expects these if ADDRESS_SANITIZER gets defined,
271  * and while ASAN does this, KASAN defines that and does not. So to avoid
272  * changing the external code, we do this.
273  */
274 #if defined(ZFS_ASAN_ENABLED)
275 #define	ADDRESS_SANITIZER 1
276 #endif
277 
278 /* Kernel space. */
279 #if defined(_KERNEL) && defined(ADDRESS_SANITIZER)
280 void __asan_unpoison_memory_region(void const volatile *addr, size_t size);
281 void __asan_poison_memory_region(void const volatile *addr, size_t size);
282 void __asan_unpoison_memory_region(void const volatile *addr, size_t size) {};
283 void __asan_poison_memory_region(void const volatile *addr, size_t size) {};
284 #endif
285 
286 /* User space. */
287 #if defined(ADDRESS_SANITIZER) && !defined(_KERNEL)
288 void __asan_unpoison_memory_region(void const volatile *addr, size_t size);
289 void __asan_poison_memory_region(void const volatile *addr, size_t size);
290 #define	ZSTD_ASAN_POISON(p, n)   __asan_poison_memory_region((p), (n))
291 #define	ZSTD_ASAN_UNPOISON(p, n) __asan_unpoison_memory_region((p), (n))
292 #else
293 #define	ZSTD_ASAN_POISON(p, n)   do { } while (0)
294 #define	ZSTD_ASAN_UNPOISON(p, n) do { } while (0)
295 #endif
296 
297 static void
298 zstd_mempool_reap(struct zstd_pool *zstd_mempool)
299 {
300 	struct zstd_pool *pool;
301 
302 	if (!zstd_mempool || !ZSTDSTAT(zstd_stat_buffers)) {
303 		return;
304 	}
305 
306 	/* free obsolete slots */
307 	for (int i = 0; i < ZSTD_POOL_MAX; i++) {
308 		pool = &zstd_mempool[i];
309 		if (pool->mem && mutex_tryenter(&pool->barrier)) {
310 			/* Free memory if unused object older than 2 minutes */
311 			if (pool->mem && gethrestime_sec() > pool->timeout) {
312 				vmem_free(pool->mem, pool->size);
313 				ZSTDSTAT_SUB(zstd_stat_buffers, 1);
314 				ZSTDSTAT_SUB(zstd_stat_size, pool->size);
315 				pool->mem = NULL;
316 				pool->size = 0;
317 				pool->timeout = 0;
318 			}
319 			mutex_exit(&pool->barrier);
320 		}
321 	}
322 }
323 
324 /*
325  * Try to get a cached allocated buffer from memory pool or allocate a new one
326  * if necessary. If a object is older than 2 minutes and does not fit the
327  * requested size, it will be released and a new cached entry will be allocated.
328  * If other pooled objects are detected without being used for 2 minutes, they
329  * will be released, too.
330  *
331  * The concept is that high frequency memory allocations of bigger objects are
332  * expensive. So if a lot of work is going on, allocations will be kept for a
333  * while and can be reused in that time frame.
334  *
335  * The scheduled release will be updated every time a object is reused.
336  */
337 
338 static void *
339 zstd_mempool_alloc(struct zstd_pool *zstd_mempool, size_t size)
340 {
341 	struct zstd_pool *pool;
342 	struct zstd_kmem *mem = NULL;
343 
344 	if (!zstd_mempool) {
345 		return (NULL);
346 	}
347 
348 	/* Seek for preallocated memory slot and free obsolete slots */
349 	for (int i = 0; i < ZSTD_POOL_MAX; i++) {
350 		pool = &zstd_mempool[i];
351 		/*
352 		 * This lock is simply a marker for a pool object being in use.
353 		 * If it's already hold, it will be skipped.
354 		 *
355 		 * We need to create it before checking it to avoid race
356 		 * conditions caused by running in a threaded context.
357 		 *
358 		 * The lock is later released by zstd_mempool_free.
359 		 */
360 		if (mutex_tryenter(&pool->barrier)) {
361 			/*
362 			 * Check if objects fits the size, if so we take it and
363 			 * update the timestamp.
364 			 */
365 			if (pool->mem && size <= pool->size) {
366 				pool->timeout = gethrestime_sec() +
367 				    ZSTD_POOL_TIMEOUT;
368 				mem = pool->mem;
369 				return (mem);
370 			}
371 			mutex_exit(&pool->barrier);
372 		}
373 	}
374 
375 	/*
376 	 * If no preallocated slot was found, try to fill in a new one.
377 	 *
378 	 * We run a similar algorithm twice here to avoid pool fragmentation.
379 	 * The first one may generate holes in the list if objects get released.
380 	 * We always make sure that these holes get filled instead of adding new
381 	 * allocations constantly at the end.
382 	 */
383 	for (int i = 0; i < ZSTD_POOL_MAX; i++) {
384 		pool = &zstd_mempool[i];
385 		if (mutex_tryenter(&pool->barrier)) {
386 			/* Object is free, try to allocate new one */
387 			if (!pool->mem) {
388 				mem = vmem_alloc(size, KM_SLEEP);
389 				if (mem) {
390 					ZSTDSTAT_ADD(zstd_stat_buffers, 1);
391 					ZSTDSTAT_ADD(zstd_stat_size, size);
392 					pool->mem = mem;
393 					pool->size = size;
394 					/* Keep track for later release */
395 					mem->pool = pool;
396 					mem->kmem_type = ZSTD_KMEM_POOL;
397 					mem->kmem_size = size;
398 				}
399 			}
400 
401 			if (size <= pool->size) {
402 				/* Update timestamp */
403 				pool->timeout = gethrestime_sec() +
404 				    ZSTD_POOL_TIMEOUT;
405 
406 				return (pool->mem);
407 			}
408 
409 			mutex_exit(&pool->barrier);
410 		}
411 	}
412 
413 	/*
414 	 * If the pool is full or the allocation failed, try lazy allocation
415 	 * instead.
416 	 */
417 	if (!mem) {
418 		mem = vmem_alloc(size, KM_NOSLEEP);
419 		if (mem) {
420 			mem->pool = NULL;
421 			mem->kmem_type = ZSTD_KMEM_DEFAULT;
422 			mem->kmem_size = size;
423 		}
424 	}
425 
426 	return (mem);
427 }
428 
429 /* Mark object as released by releasing the barrier mutex */
430 static void
431 zstd_mempool_free(struct zstd_kmem *z)
432 {
433 	/* Poison only the user-visible region (exclude header). */
434 	ZSTD_ASAN_POISON((char *)z + sizeof (struct zstd_kmem),
435 	    z->kmem_size - sizeof (struct zstd_kmem));
436 
437 	mutex_exit(&z->pool->barrier);
438 }
439 
440 /* Convert ZFS internal enum to ZSTD level */
441 static int
442 zstd_enum_to_level(enum zio_zstd_levels level, int16_t *zstd_level)
443 {
444 	if (level > 0 && level <= ZIO_ZSTD_LEVEL_19) {
445 		*zstd_level = zstd_levels[level - 1].zstd_level;
446 		return (0);
447 	}
448 	if (level >= ZIO_ZSTD_LEVEL_FAST_1 &&
449 	    level <= ZIO_ZSTD_LEVEL_FAST_1000) {
450 		*zstd_level = zstd_levels[level - ZIO_ZSTD_LEVEL_FAST_1
451 		    + ZIO_ZSTD_LEVEL_19].zstd_level;
452 		return (0);
453 	}
454 
455 	/* Invalid/unknown zfs compression enum - this should never happen. */
456 	return (1);
457 }
458 
459 #ifndef IN_LIBSA
460 /* Compress block using zstd */
461 static size_t
462 zfs_zstd_compress_impl(void *s_start, void *d_start, size_t s_len, size_t d_len,
463     int level)
464 {
465 	size_t c_len;
466 	int16_t zstd_level;
467 	zfs_zstdhdr_t *hdr;
468 	ZSTD_CCtx *cctx;
469 
470 	hdr = (zfs_zstdhdr_t *)d_start;
471 
472 	/* Skip compression if the specified level is invalid */
473 	if (zstd_enum_to_level(level, &zstd_level)) {
474 		ZSTDSTAT_BUMP(zstd_stat_com_inval);
475 		return (s_len);
476 	}
477 
478 	ASSERT3U(d_len, >=, sizeof (*hdr));
479 	ASSERT3U(d_len, <=, s_len);
480 	ASSERT3U(zstd_level, !=, 0);
481 
482 	cctx = ZSTD_createCCtx_advanced(zstd_malloc);
483 
484 	/*
485 	 * Out of kernel memory, gently fall through - this will disable
486 	 * compression in zio_compress_data
487 	 */
488 	if (!cctx) {
489 		ZSTDSTAT_BUMP(zstd_stat_com_alloc_fail);
490 		return (s_len);
491 	}
492 
493 	/* Set the compression level */
494 	ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, zstd_level);
495 
496 	/* Use the "magicless" zstd header which saves us 4 header bytes */
497 	ZSTD_CCtx_setParameter(cctx, ZSTD_c_format, ZSTD_f_zstd1_magicless);
498 
499 	/*
500 	 * Disable redundant checksum calculation and content size storage since
501 	 * this is already done by ZFS itself.
502 	 */
503 	ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 0);
504 	ZSTD_CCtx_setParameter(cctx, ZSTD_c_contentSizeFlag, 0);
505 
506 	c_len = ZSTD_compress2(cctx,
507 	    hdr->data,
508 	    d_len - sizeof (*hdr),
509 	    s_start, s_len);
510 
511 	ZSTD_freeCCtx(cctx);
512 
513 	/* Error in the compression routine, disable compression. */
514 	if (ZSTD_isError(c_len)) {
515 		/*
516 		 * If we are aborting the compression because the saves are
517 		 * too small, that is not a failure. Everything else is a
518 		 * failure, so increment the compression failure counter.
519 		 */
520 		int err = ZSTD_getErrorCode(c_len);
521 		if (err != ZSTD_error_dstSize_tooSmall) {
522 			ZSTDSTAT_BUMP(zstd_stat_com_fail);
523 			dprintf("Error: %s", ZSTD_getErrorString(err));
524 		}
525 		return (s_len);
526 	}
527 
528 	/*
529 	 * Encode the compressed buffer size at the start. We'll need this in
530 	 * decompression to counter the effects of padding which might be added
531 	 * to the compressed buffer and which, if unhandled, would confuse the
532 	 * hell out of our decompression function.
533 	 */
534 	hdr->c_len = BE_32(c_len);
535 
536 	/*
537 	 * Check version for overflow.
538 	 * The limit of 24 bits must not be exceeded. This allows a maximum
539 	 * version 1677.72.15 which we don't expect to be ever reached.
540 	 */
541 	ASSERT3U(ZSTD_VERSION_NUMBER, <=, 0xFFFFFF);
542 
543 	/*
544 	 * Encode the compression level as well. We may need to know the
545 	 * original compression level if compressed_arc is disabled, to match
546 	 * the compression settings to write this block to the L2ARC.
547 	 *
548 	 * Encode the actual level, so if the enum changes in the future, we
549 	 * will be compatible.
550 	 *
551 	 * The upper 24 bits store the ZSTD version to be able to provide
552 	 * future compatibility, since new versions might enhance the
553 	 * compression algorithm in a way, where the compressed data will
554 	 * change.
555 	 *
556 	 * As soon as such incompatibility occurs, handling code needs to be
557 	 * added, differentiating between the versions.
558 	 */
559 	zfs_set_hdrversion(hdr, ZSTD_VERSION_NUMBER);
560 	zfs_set_hdrlevel(hdr, level);
561 	hdr->raw_version_level = BE_32(hdr->raw_version_level);
562 
563 	return (c_len + sizeof (*hdr));
564 }
565 
566 static size_t
567 zfs_zstd_compress_buf(void *s_start, void *d_start, size_t s_len, size_t d_len,
568     int level)
569 {
570 	int16_t zstd_level;
571 	if (zstd_enum_to_level(level, &zstd_level)) {
572 		ZSTDSTAT_BUMP(zstd_stat_com_inval);
573 		return (s_len);
574 	}
575 	/*
576 	 * A zstd early abort heuristic.
577 	 *
578 	 * - Zeroth, if this is <= zstd-3, or < zstd_abort_size (currently
579 	 *   128k), don't try any of this, just go.
580 	 *   (because experimentally that was a reasonable cutoff for a perf win
581 	 *   with tiny ratio change)
582 	 * - First, we try LZ4 compression, and if it doesn't early abort, we
583 	 *   jump directly to whatever compression level we intended to try.
584 	 * - Second, we try zstd-1 - if that errors out (usually, but not
585 	 *   exclusively, if it would overflow), we give up early.
586 	 *
587 	 *   If it works, instead we go on and compress anyway.
588 	 *
589 	 * Why two passes? LZ4 alone gets you a lot of the way, but on highly
590 	 * compressible data, it was losing up to 8.5% of the compressed
591 	 * savings versus no early abort, and all the zstd-fast levels are
592 	 * worse indications on their own than LZ4, and don't improve the LZ4
593 	 * pass noticably if stacked like this.
594 	 */
595 	size_t actual_abort_size = zstd_abort_size;
596 	if (zstd_earlyabort_pass > 0 && zstd_level >= zstd_cutoff_level &&
597 	    s_len >= actual_abort_size) {
598 		abd_t sabd, dabd;
599 		abd_get_from_buf_struct(&sabd, s_start, s_len);
600 		abd_get_from_buf_struct(&dabd, d_start, d_len);
601 		int pass_len = zfs_lz4_compress(&sabd, &dabd, s_len, d_len, 0);
602 		abd_free(&dabd);
603 		abd_free(&sabd);
604 		if (pass_len < d_len) {
605 			ZSTDSTAT_BUMP(zstd_stat_lz4pass_allowed);
606 			goto keep_trying;
607 		}
608 		ZSTDSTAT_BUMP(zstd_stat_lz4pass_rejected);
609 
610 		pass_len = zfs_zstd_compress_impl(s_start, d_start, s_len,
611 		    d_len, ZIO_ZSTD_LEVEL_1);
612 		if (pass_len == s_len || pass_len <= 0 || pass_len > d_len) {
613 			ZSTDSTAT_BUMP(zstd_stat_zstdpass_rejected);
614 			return (s_len);
615 		}
616 		ZSTDSTAT_BUMP(zstd_stat_zstdpass_allowed);
617 	} else {
618 		ZSTDSTAT_BUMP(zstd_stat_passignored);
619 		if (s_len < actual_abort_size) {
620 			ZSTDSTAT_BUMP(zstd_stat_passignored_size);
621 		}
622 	}
623 keep_trying:
624 	return (zfs_zstd_compress_impl(s_start, d_start, s_len, d_len, level));
625 
626 }
627 #endif
628 
629 /* Decompress block using zstd and return its stored level */
630 static int
631 zfs_zstd_decompress_level_buf(void *s_start, void *d_start, size_t s_len,
632     size_t d_len, uint8_t *level)
633 {
634 	ZSTD_DCtx *dctx;
635 	size_t result;
636 	int16_t zstd_level;
637 	uint32_t c_len;
638 	const zfs_zstdhdr_t *hdr;
639 	zfs_zstdhdr_t hdr_copy;
640 
641 	hdr = (const zfs_zstdhdr_t *)s_start;
642 	c_len = BE_32(hdr->c_len);
643 
644 	/*
645 	 * Make a copy instead of directly converting the header, since we must
646 	 * not modify the original data that may be used again later.
647 	 */
648 	hdr_copy.raw_version_level = BE_32(hdr->raw_version_level);
649 	uint8_t curlevel = zfs_get_hdrlevel(&hdr_copy);
650 
651 	/*
652 	 * NOTE: We ignore the ZSTD version for now. As soon as any
653 	 * incompatibility occurs, it has to be handled accordingly.
654 	 * The version can be accessed via `hdr_copy.version`.
655 	 */
656 
657 	/*
658 	 * Convert and check the level
659 	 * An invalid level is a strong indicator for data corruption! In such
660 	 * case return an error so the upper layers can try to fix it.
661 	 */
662 	if (zstd_enum_to_level(curlevel, &zstd_level)) {
663 		ZSTDSTAT_BUMP(zstd_stat_dec_inval);
664 		return (1);
665 	}
666 
667 	ASSERT3U(d_len, >=, s_len);
668 	ASSERT3U(curlevel, !=, ZIO_COMPLEVEL_INHERIT);
669 
670 	/* Invalid compressed buffer size encoded at start */
671 	if (c_len + sizeof (*hdr) > s_len) {
672 		ZSTDSTAT_BUMP(zstd_stat_dec_header_inval);
673 		return (1);
674 	}
675 
676 	dctx = ZSTD_createDCtx_advanced(zstd_dctx_malloc);
677 	if (!dctx) {
678 		ZSTDSTAT_BUMP(zstd_stat_dec_alloc_fail);
679 		return (1);
680 	}
681 
682 	/* Set header type to "magicless" */
683 	ZSTD_DCtx_setParameter(dctx, ZSTD_d_format, ZSTD_f_zstd1_magicless);
684 
685 	/* Decompress the data and release the context */
686 	result = ZSTD_decompressDCtx(dctx, d_start, d_len, hdr->data, c_len);
687 	ZSTD_freeDCtx(dctx);
688 
689 	/*
690 	 * Returns 0 on success (decompression function returned non-negative)
691 	 * and non-zero on failure (decompression function returned negative.
692 	 */
693 	if (ZSTD_isError(result)) {
694 		ZSTDSTAT_BUMP(zstd_stat_dec_fail);
695 		return (1);
696 	}
697 
698 	/*
699 	 * An OpenZFS compressed block must expand to exactly d_len bytes.
700 	 * ZSTD_decompressDCtx returns the decompressed size on success.
701 	 */
702 	if (result != d_len) {
703 		ZSTDSTAT_BUMP(zstd_stat_dec_fail);
704 		return (1);
705 	}
706 
707 	if (level) {
708 		*level = curlevel;
709 	}
710 
711 	return (0);
712 }
713 
714 /* Decompress datablock using zstd */
715 #ifdef IN_BASE
716 int
717 zfs_zstd_decompress_buf(void *s_start, void *d_start, size_t s_len,
718     size_t d_len, int level __maybe_unused)
719 {
720 
721 	return (zfs_zstd_decompress_level_buf(s_start, d_start, s_len, d_len,
722 	    NULL));
723 }
724 #else
725 static int
726 zfs_zstd_decompress_buf(void *s_start, void *d_start, size_t s_len,
727     size_t d_len, int level __maybe_unused)
728 {
729 
730 	return (zfs_zstd_decompress_level_buf(s_start, d_start, s_len, d_len,
731 	    NULL));
732 }
733 #endif
734 
735 #ifndef IN_LIBSA
736 ZFS_COMPRESS_WRAP_DECL(zfs_zstd_compress)
737 ZFS_DECOMPRESS_WRAP_DECL(zfs_zstd_decompress)
738 ZFS_DECOMPRESS_LEVEL_WRAP_DECL(zfs_zstd_decompress_level)
739 
740 /* Allocator for zstd compression context using mempool_allocator */
741 static void *
742 zstd_alloc(void *opaque __maybe_unused, size_t size)
743 {
744 	size_t nbytes = sizeof (struct zstd_kmem) + size;
745 	struct zstd_kmem *z = NULL;
746 
747 	z = (struct zstd_kmem *)zstd_mempool_alloc(zstd_mempool_cctx, nbytes);
748 
749 	if (!z) {
750 		ZSTDSTAT_BUMP(zstd_stat_alloc_fail);
751 		return (NULL);
752 	}
753 
754 	void *p = (char *)z + sizeof (struct zstd_kmem);
755 	ZSTD_ASAN_UNPOISON(p, size);
756 	return (p);
757 }
758 
759 #endif
760 /*
761  * Allocator for zstd decompression context using mempool_allocator with
762  * fallback to reserved memory if allocation fails
763  */
764 static void *
765 zstd_dctx_alloc(void *opaque __maybe_unused, size_t size)
766 {
767 	size_t nbytes = sizeof (struct zstd_kmem) + size;
768 	struct zstd_kmem *z = NULL;
769 	enum zstd_kmem_type type = ZSTD_KMEM_DEFAULT;
770 
771 	z = (struct zstd_kmem *)zstd_mempool_alloc(zstd_mempool_dctx, nbytes);
772 	if (z) {
773 		type = ZSTD_KMEM_POOL;
774 	} else {
775 		/* Try harder, decompression shall not fail */
776 		z = vmem_alloc(nbytes, KM_SLEEP);
777 		if (z) {
778 			z->pool = NULL;
779 		}
780 		ZSTDSTAT_BUMP(zstd_stat_alloc_fail);
781 	}
782 
783 	/* Fallback if everything fails */
784 	if (!z) {
785 		/*
786 		 * Barrier since we only can handle it in a single thread. All
787 		 * other following threads need to wait here until decompression
788 		 * is completed. zstd_free will release this barrier later.
789 		 */
790 		mutex_enter(&zstd_dctx_fallback.barrier);
791 
792 		z = zstd_dctx_fallback.mem;
793 		type = ZSTD_KMEM_DCTX;
794 		ZSTDSTAT_BUMP(zstd_stat_alloc_fallback);
795 	}
796 
797 	/* Allocation should always be successful */
798 	if (!z) {
799 		return (NULL);
800 	}
801 
802 	z->kmem_type = type;
803 	z->kmem_size = nbytes;
804 
805 	void *p = (char *)z + sizeof (struct zstd_kmem);
806 	ZSTD_ASAN_UNPOISON(p, size);
807 	return (p);
808 }
809 
810 /* Free allocated memory by its specific type */
811 static void
812 zstd_free(void *opaque __maybe_unused, void *ptr)
813 {
814 	struct zstd_kmem *z =
815 	    (struct zstd_kmem *)((char *)ptr - sizeof (struct zstd_kmem));
816 	enum zstd_kmem_type type;
817 
818 	ASSERT3U(z->kmem_type, <, ZSTD_KMEM_COUNT);
819 	ASSERT3U(z->kmem_type, >, ZSTD_KMEM_UNKNOWN);
820 
821 	type = z->kmem_type;
822 	switch (type) {
823 	case ZSTD_KMEM_DEFAULT:
824 		vmem_free(z, z->kmem_size);
825 		break;
826 	case ZSTD_KMEM_POOL:
827 		zstd_mempool_free(z);
828 		break;
829 	case ZSTD_KMEM_DCTX:
830 		/* Poison fallback user region on release. */
831 		ZSTD_ASAN_POISON(ptr, z->kmem_size - sizeof (struct zstd_kmem));
832 		mutex_exit(&zstd_dctx_fallback.barrier);
833 		break;
834 	default:
835 		break;
836 	}
837 }
838 
839 /* Allocate fallback memory to ensure safe decompression */
840 static void __init
841 create_fallback_mem(struct zstd_fallback_mem *mem, size_t size)
842 {
843 	mem->mem_size = size;
844 	mem->mem = vmem_zalloc(mem->mem_size, KM_SLEEP);
845 	mutex_init(&mem->barrier, NULL, MUTEX_DEFAULT, NULL);
846 }
847 
848 /* Initialize memory pool barrier mutexes */
849 static void __init
850 zstd_mempool_init(void)
851 {
852 	zstd_mempool_cctx =
853 	    vmem_zalloc(ZSTD_POOL_MAX * sizeof (struct zstd_pool), KM_SLEEP);
854 	zstd_mempool_dctx =
855 	    vmem_zalloc(ZSTD_POOL_MAX * sizeof (struct zstd_pool), KM_SLEEP);
856 
857 	for (int i = 0; i < ZSTD_POOL_MAX; i++) {
858 		mutex_init(&zstd_mempool_cctx[i].barrier, NULL,
859 		    MUTEX_DEFAULT, NULL);
860 		mutex_init(&zstd_mempool_dctx[i].barrier, NULL,
861 		    MUTEX_DEFAULT, NULL);
862 	}
863 }
864 
865 /* Initialize zstd-related memory handling */
866 static int __init
867 zstd_meminit(void)
868 {
869 	zstd_mempool_init();
870 
871 	/*
872 	 * Estimate the size of the fallback decompression context.
873 	 * The expected size on x64 with current ZSTD should be about 160 KB.
874 	 */
875 	create_fallback_mem(&zstd_dctx_fallback,
876 	    P2ROUNDUP(ZSTD_estimateDCtxSize() + sizeof (struct zstd_kmem),
877 	    PAGESIZE));
878 
879 	return (0);
880 }
881 
882 /* Release object from pool and free memory */
883 static void
884 release_pool(struct zstd_pool *pool)
885 {
886 	mutex_destroy(&pool->barrier);
887 	vmem_free(pool->mem, pool->size);
888 	pool->mem = NULL;
889 	pool->size = 0;
890 }
891 
892 /* Release memory pool objects */
893 static void
894 zstd_mempool_deinit(void)
895 {
896 	for (int i = 0; i < ZSTD_POOL_MAX; i++) {
897 		release_pool(&zstd_mempool_cctx[i]);
898 		release_pool(&zstd_mempool_dctx[i]);
899 	}
900 
901 	vmem_free(zstd_mempool_dctx, ZSTD_POOL_MAX * sizeof (struct zstd_pool));
902 	vmem_free(zstd_mempool_cctx, ZSTD_POOL_MAX * sizeof (struct zstd_pool));
903 	zstd_mempool_dctx = NULL;
904 	zstd_mempool_cctx = NULL;
905 }
906 
907 /* release unused memory from pool */
908 
909 void
910 zfs_zstd_cache_reap_now(void)
911 {
912 
913 	/*
914 	 * Short-circuit if there are no buffers to begin with.
915 	 */
916 	if (ZSTDSTAT(zstd_stat_buffers) == 0)
917 		return;
918 
919 	/*
920 	 * calling alloc with zero size seeks
921 	 * and releases old unused objects
922 	 */
923 	zstd_mempool_reap(zstd_mempool_cctx);
924 	zstd_mempool_reap(zstd_mempool_dctx);
925 }
926 
927 extern int __init
928 zstd_init(void)
929 {
930 	/* Set pool size by using maximum sane thread count * 4 */
931 	pool_count = (boot_ncpus * 4);
932 	zstd_meminit();
933 
934 	/* Initialize kstat */
935 	zstd_ksp = kstat_create("zfs", 0, "zstd", "misc",
936 	    KSTAT_TYPE_NAMED, sizeof (zstd_stats) / sizeof (kstat_named_t),
937 	    KSTAT_FLAG_VIRTUAL);
938 	if (zstd_ksp != NULL) {
939 		zstd_ksp->ks_data = &zstd_stats;
940 		kstat_install(zstd_ksp);
941 #ifdef _KERNEL
942 		zstd_ksp->ks_update = kstat_zstd_update;
943 #endif
944 	}
945 
946 	return (0);
947 }
948 
949 extern void
950 zstd_fini(void)
951 {
952 	/* Deinitialize kstat */
953 	if (zstd_ksp != NULL) {
954 		kstat_delete(zstd_ksp);
955 		zstd_ksp = NULL;
956 	}
957 
958 	/* Release fallback memory */
959 	vmem_free(zstd_dctx_fallback.mem, zstd_dctx_fallback.mem_size);
960 	mutex_destroy(&zstd_dctx_fallback.barrier);
961 
962 	/* Deinit memory pool */
963 	zstd_mempool_deinit();
964 }
965 
966 #if defined(_KERNEL)
967 #ifdef __FreeBSD__
968 module_init(zstd_init);
969 module_exit(zstd_fini);
970 #endif
971 
972 ZFS_MODULE_PARAM(zfs, zstd_, earlyabort_pass, UINT, ZMOD_RW,
973 	"Enable early abort attempts when using zstd");
974 ZFS_MODULE_PARAM(zfs, zstd_, abort_size, UINT, ZMOD_RW,
975 	"Minimal size of block to attempt early abort");
976 #endif
977