compress: wait for an abandoned frame's jobs before freeing what they use

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;
     }

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);

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   === */

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.

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);