compress: fix ZSTD_CCtx_refThreadPool() rebinding and shared pool resizing

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

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

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;

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