diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 88e1df51232..c51fd1853ff 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1360,6 +1360,15 @@ size_t ZSTD_CCtx_refThreadPool(ZSTD_CCtx* cctx, ZSTD_threadPool* pool) { RETURN_ERROR_IF(cctx->streamStage != zcss_init, stage_wrong, "Can't ref a pool when ctx not in init stage."); +#ifdef ZSTD_MULTITHREAD + if (pool != cctx->pool) { + /* The MT context, created by the first multi-threaded frame and then reused, + * runs its jobs on the pool it was created with : without this, the next frames + * would still use the previous pool, which the caller may free. No job is running : + * the stage is zcss_init, and ZSTDMT_freeCCtx() waits for pending jobs. */ + ZSTDMT_freeCCtx(cctx->mtctx); cctx->mtctx = NULL; + } +#endif cctx->pool = pool; return 0; } diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 7118bbe7f4a..38ad31d0e0d 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -1091,7 +1091,9 @@ size_t ZSTDMT_sizeof_CCtx(ZSTDMT_CCtx* mtctx) * @return : error code if fails, 0 on success */ static size_t ZSTDMT_resize(ZSTDMT_CCtx* mtctx, unsigned nbWorkers) { - if (POOL_resize(mtctx->factory, nbWorkers)) return ERROR(memory_allocation); + if (!mtctx->providedFactory) { /* a pool provided by the application keeps its size */ + if (POOL_resize(mtctx->factory, nbWorkers)) return ERROR(memory_allocation); + } FORWARD_IF_ERROR( ZSTDMT_expandJobsTable(mtctx, nbWorkers) , ""); mtctx->bufPool = ZSTDMT_expandBufferPool(mtctx->bufPool, BUF_POOL_MAX_NB_BUFFERS(nbWorkers)); if (mtctx->bufPool == NULL) return ERROR(memory_allocation); diff --git a/lib/zstd.h b/lib/zstd.h index d06c05edd80..e0862eb3fa8 100644 --- a/lib/zstd.h +++ b/lib/zstd.h @@ -1899,9 +1899,10 @@ ZSTDLIB_STATIC_API ZSTD_CDict* ZSTD_createCDict_advanced(const void* dict, size_ * This can limit resources for applications with multiple threads where each one uses * a threaded compression mode (via ZSTD_c_nbWorkers parameter). * ZSTD_createThreadPool creates a new thread pool with a given number of threads. + * ZSTD_c_nbWorkers does not change the number of threads of such a pool. * Note that the lifetime of such pool must exist while being used. * ZSTD_CCtx_refThreadPool assigns a thread pool to a context (use NULL argument value - * to use an internal thread pool). + * to use an internal thread pool). The context then no longer uses its previous pool. * ZSTD_freeThreadPool frees a thread pool, accepts NULL pointer. */ typedef struct POOL_ctx_s ZSTD_threadPool; diff --git a/tests/fuzzer.c b/tests/fuzzer.c index 3d3cd66abf8..cbc324a98ca 100644 --- a/tests/fuzzer.c +++ b/tests/fuzzer.c @@ -44,6 +44,7 @@ /* must be included after util.h, due to ERROR macro redefinition issue on Visual Studio */ #include "zstd_internal.h" /* ZSTD_WORKSPACETOOLARGE_MAXDURATION, ZSTD_WORKSPACETOOLARGE_FACTOR, KB, MB */ #include "threading.h" /* ZSTD_pthread_create, ZSTD_pthread_join */ +#include "pool.h" /* POOL_sizeof */ #include "compress/hist.h" /* HIST_count_wksp */ #include "compress/zstd_compress_internal.h" /* ZSTD_get1BlockSummary */ @@ -483,6 +484,31 @@ static int threadPoolTests(void) { } DISPLAYLEVEL(3, "OK \n"); + DISPLAYLEVEL(3, "thread pool test : change the pool and nbWorkers of a context: "); + { ZSTD_CCtx* const cctx = ZSTD_createCCtx(); + ZSTD_threadPool* const pool1 = ZSTD_createThreadPool(1); + ZSTD_threadPool* const pool2 = ZSTD_createThreadPool(1); + CHECK(cctx != NULL && pool1 != NULL && pool2 != NULL); + CHECK(!ZSTD_isError( ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, 2) )); + CHECK(!ZSTD_isError( ZSTD_CCtx_refThreadPool(cctx, pool1) )); + CHECK(!ZSTD_isError( ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, CNBuffSize) )); + /* the next frames must use pool2, then an internal pool */ + CHECK(!ZSTD_isError( ZSTD_CCtx_refThreadPool(cctx, pool2) )); + ZSTD_freeThreadPool(pool1); + CHECK(!ZSTD_isError( ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, CNBuffSize) )); + { size_t const pool2Size = POOL_sizeof(pool2); + /* the context must not resize a pool it shares */ + CHECK(!ZSTD_isError( ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, 4) )); + CHECK(!ZSTD_isError( ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, CNBuffSize) )); + CHECK(POOL_sizeof(pool2) == pool2Size); + } + CHECK(!ZSTD_isError( ZSTD_CCtx_refThreadPool(cctx, NULL) )); + ZSTD_freeThreadPool(pool2); + CHECK(!ZSTD_isError( ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, CNBuffSize) )); + ZSTD_freeCCtx(cctx); + } + DISPLAYLEVEL(3, "OK \n"); + _end: free(CNBuffer); free(compressedBuffer);