From ed26f948cd553ed4f35e681c06bb03be3426f2b8 Mon Sep 17 00:00:00 2001 From: Eshan05 Date: Wed, 28 Oct 2026 21:41:14 +0530 Subject: [PATCH] compress: wait for an abandoned frame's jobs before freeing what they use With nbWorkers >= 1, a job of an abandoned frame can still be running when the context frees memory it uses: - ZSTD_freeCCtx() frees the dictionary before ZSTDMT_freeCCtx() joins the workers; - ZSTD_CCtx_reset(session_and_parameters), or a session reset followed by loading another dictionary, frees it while the job runs; - with ZSTD_CCtx_refThreadPool(), ZSTDMT_freeCCtx() does not stop the shared pool and frees the jobs table and CCtx pool under the worker; - after a session reset, raising nbWorkers resizes the CCtx pool before the abandoned jobs are waited for. Move the "previous compression not correctly finished" wait of ZSTDMT_initCStream_internal() into ZSTDMT_waitForPendingJobs(), and call it on session reset, at the start of ZSTDMT_freeCCtx(), and before ZSTDMT_resize(). Free the MT context before the dictionaries. Fixes #4862 Signed-off-by: Eshan05 --- lib/compress/zstd_compress.c | 8 ++++-- lib/compress/zstdmt_compress.c | 15 +++++++---- lib/compress/zstdmt_compress.h | 6 +++++ lib/zstd.h | 1 + tests/zstreamtest.c | 47 ++++++++++++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 7 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 38c3fc0d4cd..88e1df51232 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -184,10 +184,10 @@ static void ZSTD_freeCCtxContent(ZSTD_CCtx* cctx) { assert(cctx != NULL); assert(cctx->staticSize == 0); - ZSTD_clearAllDicts(cctx); #ifdef ZSTD_MULTITHREAD - ZSTDMT_freeCCtx(cctx->mtctx); cctx->mtctx = NULL; + ZSTDMT_freeCCtx(cctx->mtctx); cctx->mtctx = NULL; /* pending jobs may use the dictionary */ #endif + ZSTD_clearAllDicts(cctx); ZSTD_cwksp_free(&cctx->workspace, cctx->customMem); } @@ -1389,6 +1389,10 @@ size_t ZSTD_CCtx_reset(ZSTD_CCtx* cctx, ZSTD_ResetDirective reset) { if ( (reset == ZSTD_reset_session_only) || (reset == ZSTD_reset_session_and_parameters) ) { +#ifdef ZSTD_MULTITHREAD + /* jobs of an abandoned frame may still reference the dictionary */ + if (cctx->mtctx != NULL) ZSTDMT_waitForPendingJobs(cctx->mtctx); +#endif cctx->streamStage = zcss_init; cctx->pledgedSrcSizePlusOne = 0; } diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index cb710475c2a..1594221c4e9 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -1043,9 +1043,18 @@ static void ZSTDMT_waitForAllJobsCompleted(ZSTDMT_CCtx* mtctx) } } +void ZSTDMT_waitForPendingJobs(ZSTDMT_CCtx* mtctx) +{ + if (mtctx->allJobsCompleted == 0) { /* previous compression not correctly finished */ + ZSTDMT_waitForAllJobsCompleted(mtctx); + ZSTDMT_releaseAllJobResources(mtctx); /* Will set allJobsCompleted to 1 */ + } +} + size_t ZSTDMT_freeCCtx(ZSTDMT_CCtx* mtctx) { if (mtctx==NULL) return 0; /* compatible with free on NULL */ + ZSTDMT_waitForPendingJobs(mtctx); /* a shared pool is not stopped below */ if (!mtctx->providedFactory) POOL_free(mtctx->factory); /* stop and free worker threads */ ZSTDMT_releaseAllJobResources(mtctx); /* release job resources into pools first */ @@ -1262,17 +1271,13 @@ size_t ZSTDMT_initCStream_internal( assert(!((dict) && (cdict))); /* either dict or cdict, not both */ /* init */ + ZSTDMT_waitForPendingJobs(mtctx); /* before resizing the pools they use */ if (params.nbWorkers != mtctx->params.nbWorkers) FORWARD_IF_ERROR( ZSTDMT_resize(mtctx, (unsigned)params.nbWorkers) , ""); if (params.jobSize != 0 && params.jobSize < ZSTDMT_JOBSIZE_MIN) params.jobSize = ZSTDMT_JOBSIZE_MIN; if (params.jobSize > (size_t)ZSTDMT_JOBSIZE_MAX) params.jobSize = (size_t)ZSTDMT_JOBSIZE_MAX; - if (mtctx->allJobsCompleted == 0) { /* previous compression not correctly finished */ - ZSTDMT_waitForAllJobsCompleted(mtctx); - ZSTDMT_releaseAllJobResources(mtctx); /* Will set allJobsCompleted to 1 */ - } - mtctx->params = params; mtctx->frameContentSize = pledgedSrcSize; ZSTD_freeCDict(mtctx->cdictLocal); diff --git a/lib/compress/zstdmt_compress.h b/lib/compress/zstdmt_compress.h index 91b489b9cb4..3e2ee664ed0 100644 --- a/lib/compress/zstdmt_compress.h +++ b/lib/compress/zstdmt_compress.h @@ -49,6 +49,12 @@ ZSTDMT_CCtx* ZSTDMT_createCCtx_advanced(unsigned nbWorkers, ZSTD_threadPool *pool); size_t ZSTDMT_freeCCtx(ZSTDMT_CCtx* mtctx); +/*! ZSTDMT_waitForPendingJobs() : + * Waits for the jobs of an unfinished frame to complete, then releases their resources. + * Must be called before freeing anything these jobs reference, such as the dictionary. + * Does nothing if the last frame was completed. */ +void ZSTDMT_waitForPendingJobs(ZSTDMT_CCtx* mtctx); + size_t ZSTDMT_sizeof_CCtx(ZSTDMT_CCtx* mtctx); /* === Streaming functions === */ diff --git a/lib/zstd.h b/lib/zstd.h index 97fef316fd6..3c24acc230d 100644 --- a/lib/zstd.h +++ b/lib/zstd.h @@ -597,6 +597,7 @@ typedef enum { * - The session : will stop compressing current frame, and make CCtx ready to start a new one. * Useful after an error, or to interrupt any ongoing compression. * Any internal data not yet flushed is cancelled. + * With nbWorkers >= 1, it waits for the jobs already started to complete. * Compression parameters and dictionary remain unchanged. * They will be used to compress next frame. * Resetting session never fails. diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c index 88a20a3e270..22b038828cd 100644 --- a/tests/zstreamtest.c +++ b/tests/zstreamtest.c @@ -2512,6 +2512,53 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests) } DISPLAYLEVEL(3, "OK \n"); +#ifdef ZSTD_MULTITHREAD + /* Freeing or resetting the context must wait for the jobs still running, + * since they use its dictionary, pools and jobs table */ + DISPLAYLEVEL(3, "test%3i : abandon a multi-threaded frame with a job in progress : ", testNb++); + { size_t const jobSize = 1 MB; + size_t const srcSize = jobSize + 64 KB; /* starts one job, the rest stays buffered */ + size_t const largeDictSize = 1 MB; + const void* const largeDict = (const char*)CNBuffer + srcSize; + int scenario; + for (scenario = 0; scenario < 4; scenario++) { + ZSTD_CCtx* const cctx = ZSTD_createCCtx(); + ZSTD_threadPool* const pool = (scenario == 1) ? ZSTD_createThreadPool(1) : NULL; + ZSTD_inBuffer in = { CNBuffer, srcSize, 0 }; + ZSTD_outBuffer out = { compressedBuffer, compressedBufferSize, 0 }; + size_t nextDictSize = largeDictSize; + CHECK(cctx == NULL, "ZSTD_createCCtx failed"); + CHECK_Z( ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, 10) ); + CHECK_Z( ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, 1) ); + CHECK_Z( ZSTD_CCtx_setParameter(cctx, ZSTD_c_jobSize, (int)jobSize) ); + if (pool != NULL) CHECK_Z( ZSTD_CCtx_refThreadPool(cctx, pool) ); + CHECK_Z( ZSTD_CCtx_loadDictionary(cctx, largeDict, largeDictSize) ); + CHECK_Z( ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue) ); + CHECK(in.pos < jobSize, "the first job should be in progress"); + if (scenario >= 2) { + /* abandon the frame, then change the dictionary or the number of workers */ + CHECK_Z( ZSTD_CCtx_reset(cctx, ZSTD_reset_session_only) ); + if (scenario == 2) { + nextDictSize = largeDictSize / 2; + CHECK_Z( ZSTD_CCtx_loadDictionary(cctx, largeDict, nextDictSize) ); + } else { + CHECK_Z( ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, 4) ); + } + cSize = ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, srcSize); + CHECK_Z(cSize); + { size_t const dSize = ZSTD_decompress_usingDict(zd, decodedBuffer, decodedBufferSize, + compressedBuffer, cSize, largeDict, nextDictSize); + CHECK_Z(dSize); + CHECK(dSize != srcSize, "wrong decompressed size"); + CHECK(memcmp(decodedBuffer, CNBuffer, srcSize), "decompressed data must match"); + } } + ZSTD_freeCCtx(cctx); /* scenarios 0 and 1 free it with the job in progress */ + ZSTD_freeThreadPool(pool); + } + } + DISPLAYLEVEL(3, "OK \n"); +#endif + _end: FUZ_freeDictionary(dictionary); ZSTD_freeCStream(zc);