zstdmt: give the thread pools a queue so small jobs run in parallel

lib/common/pool.c

@@ -103,9 +103,11 @@ static void* POOL_thread(void* opaque) {
     assert(0);  /* Unreachable */
 }
 
-/* ZSTD_createThreadPool() : public access point */
+/* ZSTD_createThreadPool() : public access point
+ * The queue lets a compression context post a job without waiting for a worker
+ * to wake up and take the previous one (see ZSTDMT_POOL_QUEUE_SIZE). */
 POOL_ctx* ZSTD_createThreadPool(size_t numThreads) {
-    return POOL_create (numThreads, 0);
+    return POOL_create (numThreads, numThreads);
 }
 
 POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {

lib/compress/zstdmt_compress.c

@@ -281,6 +281,15 @@ static void ZSTDMT_releaseBuffer(ZSTDMT_bufferPool* bufPool, Buffer buf)
  * So we only need one seq buffer per worker. */
 #define SEQ_POOL_MAX_NB_BUFFERS(nbWorkers) (nbWorkers)
 
+/* Queue of the internal thread pool.
+ * Without a queue, POOL_tryAdd() fails until a worker has taken the previous job,
+ * so a job prepared faster than a worker wakes up has to wait for the oldest job to finish,
+ * and small jobs end up running one after another.
+ * A queued job doesn't hold any buffer : its input is already in the round buffer,
+ * and the worker gets an output buffer when it starts the job.
+ * The number of jobs in flight is still bounded by the jobs table and the round buffer. */
+#define ZSTDMT_POOL_QUEUE_SIZE(nbWorkers) (nbWorkers)
+
 /* =====   Seq Pool Wrapper   ====== */
 
 typedef ZSTDMT_bufferPool ZSTDMT_seqPool;
@@ -874,7 +883,7 @@ struct ZSTDMT_CCtx_s {
     ZSTD_CCtx_params params;
     size_t targetSectionSize;
     size_t targetPrefixSize;
-    int jobReady;        /* 1 => one job is already prepared, but pool has shortage of workers. Don't create a new job. */
+    int jobReady;        /* 1 => one job is already prepared, but the pool could not accept it. Don't create a new job. */
     InBuff_t inBuff;
     RoundBuff_t roundBuff;
     SerialState serial;
@@ -890,6 +899,7 @@ struct ZSTDMT_CCtx_s {
     ZSTD_customMem cMem;
     ZSTD_CDict* cdictLocal;
     const ZSTD_CDict* cdict;
+    unsigned factoryQueueSize;   /* queue size of the internal pool, unused with a provided pool */
     unsigned providedFactory: 1;
 };
 
@@ -972,7 +982,8 @@ MEM_STATIC ZSTDMT_CCtx* ZSTDMT_createCCtx_advanced_internal(unsigned nbWorkers,
       mtctx->providedFactory = 1;
     }
     else {
-      mtctx->factory = POOL_create_advanced(nbWorkers, 0, cMem);
+      mtctx->factory = POOL_create_advanced(nbWorkers, ZSTDMT_POOL_QUEUE_SIZE(nbWorkers), cMem);
+      mtctx->factoryQueueSize = ZSTDMT_POOL_QUEUE_SIZE(nbWorkers);
       mtctx->providedFactory = 0;
     }
     mtctx->jobs = ZSTDMT_createJobsTable(&nbJobs, cMem);
@@ -1092,7 +1103,15 @@ size_t ZSTDMT_sizeof_CCtx(ZSTDMT_CCtx* mtctx)
 static size_t ZSTDMT_resize(ZSTDMT_CCtx* mtctx, unsigned nbWorkers)
 {
     if (!mtctx->providedFactory) {   /* a pool provided by the application keeps its size */
-        if (POOL_resize(mtctx->factory, nbWorkers)) return ERROR(memory_allocation);
+        if (ZSTDMT_POOL_QUEUE_SIZE(nbWorkers) > mtctx->factoryQueueSize) {
+            /* POOL_resize() can't enlarge the queue, so replace the internal pool.
+             * Its workers are idle : all jobs completed before resizing. */
+            POOL_ctx* const factory = POOL_create_advanced(nbWorkers, ZSTDMT_POOL_QUEUE_SIZE(nbWorkers), mtctx->cMem);
+            if (factory == NULL) return ERROR(memory_allocation);
+            POOL_free(mtctx->factory);
+            mtctx->factory = factory;
+            mtctx->factoryQueueSize = ZSTDMT_POOL_QUEUE_SIZE(nbWorkers);
+        } else 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));

tests/zstreamtest.c

@@ -2538,8 +2538,11 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests)
                 CHECK_Z( ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end) );
             }
             if (scenario == 4) {
-                /* the second job is prepared, and waits until the worker is free */
-                in.size = 2 * jobSize;
+                /* once the worker has started the first job, the second job waits in the pool queue,
+                 * and the third one is prepared, and waits until the queue has room */
+                while (ZSTD_getFrameProgression(cctx).consumed == 0) {}
+                in.size = 3 * jobSize;
+                CHECK_Z( ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue) );
                 CHECK_Z( ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue) );
             }
             if (scenario >= 3) {
@@ -2553,6 +2556,45 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests)
         }
     }
     DISPLAYLEVEL(3, "OK \n");
+
+    /* The pool must accept a new job while all its workers are busy,
+     * otherwise the next job waits for the oldest one to finish,
+     * and small jobs run one after another (#4890) */
+    DISPLAYLEVEL(3, "test%3i : a new job is queued while all workers are busy : ", testNb++);
+    {   size_t const jobSize = 2 MB;
+        int scenario;
+        for (scenario = 0; scenario < 2; scenario++) {
+            ZSTD_CCtx* const cctx = ZSTD_createCCtx();
+            ZSTD_threadPool* const pool = (scenario == 1) ? ZSTD_createThreadPool(1) : NULL;
+            ZSTD_inBuffer in = { CNBuffer, jobSize, 0 };
+            ZSTD_outBuffer out = { compressedBuffer, compressedBufferSize, 0 };
+            size_t remaining;
+            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_compressStream2(cctx, &out, &in, ZSTD_e_continue) );   /* posts the first job */
+            /* the worker reports progress once it has compressed the first part of the job,
+             * and it is still busy with the rest of it */
+            while (ZSTD_getFrameProgression(cctx).consumed == 0) {}
+            in.size = 2 * jobSize;
+            CHECK_Z( ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue) );   /* posts the second job */
+            CHECK(ZSTD_getFrameProgression(cctx).currentJobID != 2, "the second job was not accepted by the pool");
+            do {
+                remaining = ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end);
+                CHECK_Z(remaining);
+            } while (remaining != 0);
+            {   size_t const dSize = ZSTD_decompress(decodedBuffer, decodedBufferSize, out.dst, out.pos);
+                CHECK_Z(dSize);
+                CHECK(dSize != 2 * jobSize, "wrong regenerated size");
+                CHECK(memcmp(decodedBuffer, CNBuffer, dSize) != 0, "corruption");
+            }
+            ZSTD_freeCCtx(cctx);
+            ZSTD_freeThreadPool(pool);
+        }
+    }
+    DISPLAYLEVEL(3, "OK \n");
 #endif
 
 _end: