🌐 AI搜索 & 代理 主页
blob: 943c20aba5e79dcb95694f53725a1a6fc99b82af [file] [log] [blame]
Francois Visconte77df43e2021-11-18 12:49:401#ifndef USE_EXTERNAL_ZSTD
ghatdevc3115f62018-01-25 16:01:312/*
Vianney Tranca4d3c72023-04-20 20:30:323 * Copyright (c) Meta Platforms, Inc. and affiliates.
Vianney Tran45f76132017-02-22 00:06:474 * All rights reserved.
5 *
ghatdevc3115f62018-01-25 16:01:316 * This source code is licensed under both the BSD-style license (found in the
7 * LICENSE file in the root directory of this source tree) and the GPLv2 (found
8 * in the COPYING file in the root directory of this source tree).
9 * You may select, at your option, one of the above-listed licenses.
Vianney Tran45f76132017-02-22 00:06:4710 */
11
12
13/* ====== Dependencies ======= */
Vianney Tranca4d3c72023-04-20 20:30:3214#include "allocations.h" /* ZSTD_customCalloc, ZSTD_customFree */
Vianney Tran9f7d8ae2021-02-04 18:32:4815#include "zstd_deps.h" /* size_t */
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3416#include "debug.h" /* assert */
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3417#include "pool.h"
Vianney Tran45f76132017-02-22 00:06:4718
19/* ====== Compiler specifics ====== */
20#if defined(_MSC_VER)
21# pragma warning(disable : 4204) /* disable: C4204: non-constant aggregate initializer */
22#endif
23
24
25#ifdef ZSTD_MULTITHREAD
26
27#include "threading.h" /* pthread adaptation */
28
29/* A job is a function and an opaque argument */
30typedef struct POOL_job_s {
ghatdevc3115f62018-01-25 16:01:3131 POOL_function function;
32 void *opaque;
Vianney Tran45f76132017-02-22 00:06:4733} POOL_job;
34
35struct POOL_ctx_s {
ghatdevc3115f62018-01-25 16:01:3136 ZSTD_customMem customMem;
Vianney Tran45f76132017-02-22 00:06:4737 /* Keep track of the threads */
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3438 ZSTD_pthread_t* threads;
39 size_t threadCapacity;
40 size_t threadLimit;
Vianney Tran45f76132017-02-22 00:06:4741
42 /* The queue is a circular buffer */
43 POOL_job *queue;
44 size_t queueHead;
45 size_t queueTail;
46 size_t queueSize;
ghatdevc3115f62018-01-25 16:01:3147
48 /* The number of threads working on jobs */
49 size_t numThreadsBusy;
50 /* Indicates if the queue is empty */
51 int queueEmpty;
52
Vianney Tran45f76132017-02-22 00:06:4753 /* The mutex protects the queue */
ghatdevc3115f62018-01-25 16:01:3154 ZSTD_pthread_mutex_t queueMutex;
Vianney Tran45f76132017-02-22 00:06:4755 /* Condition variable for pushers to wait on when the queue is full */
ghatdevc3115f62018-01-25 16:01:3156 ZSTD_pthread_cond_t queuePushCond;
Vianney Tran45f76132017-02-22 00:06:4757 /* Condition variables for poppers to wait on when the queue is empty */
ghatdevc3115f62018-01-25 16:01:3158 ZSTD_pthread_cond_t queuePopCond;
Vianney Tran45f76132017-02-22 00:06:4759 /* Indicates if the queue is shutting down */
60 int shutdown;
61};
62
63/* POOL_thread() :
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3464 * Work thread for the thread pool.
65 * Waits for jobs and executes them.
66 * @returns : NULL on failure else non-null.
67 */
Vianney Tran45f76132017-02-22 00:06:4768static void* POOL_thread(void* opaque) {
69 POOL_ctx* const ctx = (POOL_ctx*)opaque;
70 if (!ctx) { return NULL; }
71 for (;;) {
72 /* Lock the mutex and wait for a non-empty queue or until shutdown */
ghatdevc3115f62018-01-25 16:01:3173 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
74
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3475 while ( ctx->queueEmpty
76 || (ctx->numThreadsBusy >= ctx->threadLimit) ) {
77 if (ctx->shutdown) {
78 /* even if !queueEmpty, (possible if numThreadsBusy >= threadLimit),
79 * a few threads will be shutdown while !queueEmpty,
80 * but enough threads will remain active to finish the queue */
81 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
82 return opaque;
83 }
ghatdevc3115f62018-01-25 16:01:3184 ZSTD_pthread_cond_wait(&ctx->queuePopCond, &ctx->queueMutex);
Vianney Tran45f76132017-02-22 00:06:4785 }
Vianney Tran45f76132017-02-22 00:06:4786 /* Pop a job off the queue */
87 { POOL_job const job = ctx->queue[ctx->queueHead];
88 ctx->queueHead = (ctx->queueHead + 1) % ctx->queueSize;
ghatdevc3115f62018-01-25 16:01:3189 ctx->numThreadsBusy++;
Vianney Tran7c9ad242022-04-14 16:08:3090 ctx->queueEmpty = (ctx->queueHead == ctx->queueTail);
Vianney Tran45f76132017-02-22 00:06:4791 /* Unlock the mutex, signal a pusher, and run the job */
ghatdevc3115f62018-01-25 16:01:3192 ZSTD_pthread_cond_signal(&ctx->queuePushCond);
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3493 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
ghatdevc3115f62018-01-25 16:01:3194
Vianney Tran45f76132017-02-22 00:06:4795 job.function(job.opaque);
ghatdevc3115f62018-01-25 16:01:3196
97 /* If the intended queue size was 0, signal after finishing job */
Vlad Gorodetskyf38c0ed2019-04-08 11:45:3498 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
99 ctx->numThreadsBusy--;
Vianney Tranca4d3c72023-04-20 20:30:32100 ZSTD_pthread_cond_signal(&ctx->queuePushCond);
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34101 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
102 }
ghatdevc3115f62018-01-25 16:01:31103 } /* for (;;) */
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34104 assert(0); /* Unreachable */
Vianney Tran45f76132017-02-22 00:06:47105}
106
Vianney Tran7c9ad242022-04-14 16:08:30107/* ZSTD_createThreadPool() : public access point */
Vianney Tran9f7d8ae2021-02-04 18:32:48108POOL_ctx* ZSTD_createThreadPool(size_t numThreads) {
109 return POOL_create (numThreads, 0);
110}
111
ghatdevc3115f62018-01-25 16:01:31112POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
113 return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
114}
115
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34116POOL_ctx* POOL_create_advanced(size_t numThreads, size_t queueSize,
Vianney Tran7c9ad242022-04-14 16:08:30117 ZSTD_customMem customMem)
118{
ghatdevc3115f62018-01-25 16:01:31119 POOL_ctx* ctx;
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34120 /* Check parameters */
ghatdevc3115f62018-01-25 16:01:31121 if (!numThreads) { return NULL; }
Vianney Tran45f76132017-02-22 00:06:47122 /* Allocate the context and zero initialize */
Vianney Tran9f7d8ae2021-02-04 18:32:48123 ctx = (POOL_ctx*)ZSTD_customCalloc(sizeof(POOL_ctx), customMem);
Vianney Tran45f76132017-02-22 00:06:47124 if (!ctx) { return NULL; }
125 /* Initialize the job queue.
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34126 * It needs one extra space since one space is wasted to differentiate
127 * empty and full queues.
Vianney Tran45f76132017-02-22 00:06:47128 */
129 ctx->queueSize = queueSize + 1;
Vianney Tranca4d3c72023-04-20 20:30:32130 ctx->queue = (POOL_job*)ZSTD_customCalloc(ctx->queueSize * sizeof(POOL_job), customMem);
Vianney Tran45f76132017-02-22 00:06:47131 ctx->queueHead = 0;
132 ctx->queueTail = 0;
ghatdevc3115f62018-01-25 16:01:31133 ctx->numThreadsBusy = 0;
134 ctx->queueEmpty = 1;
Dan Benamyd5561582019-11-06 03:19:47135 {
136 int error = 0;
137 error |= ZSTD_pthread_mutex_init(&ctx->queueMutex, NULL);
138 error |= ZSTD_pthread_cond_init(&ctx->queuePushCond, NULL);
139 error |= ZSTD_pthread_cond_init(&ctx->queuePopCond, NULL);
140 if (error) { POOL_free(ctx); return NULL; }
141 }
Vianney Tran45f76132017-02-22 00:06:47142 ctx->shutdown = 0;
143 /* Allocate space for the thread handles */
Vianney Tranca4d3c72023-04-20 20:30:32144 ctx->threads = (ZSTD_pthread_t*)ZSTD_customCalloc(numThreads * sizeof(ZSTD_pthread_t), customMem);
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34145 ctx->threadCapacity = 0;
ghatdevc3115f62018-01-25 16:01:31146 ctx->customMem = customMem;
Vianney Tran45f76132017-02-22 00:06:47147 /* Check for errors */
148 if (!ctx->threads || !ctx->queue) { POOL_free(ctx); return NULL; }
149 /* Initialize the threads */
150 { size_t i;
151 for (i = 0; i < numThreads; ++i) {
ghatdevc3115f62018-01-25 16:01:31152 if (ZSTD_pthread_create(&ctx->threads[i], NULL, &POOL_thread, ctx)) {
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34153 ctx->threadCapacity = i;
Vianney Tran45f76132017-02-22 00:06:47154 POOL_free(ctx);
155 return NULL;
156 } }
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34157 ctx->threadCapacity = numThreads;
158 ctx->threadLimit = numThreads;
Vianney Tran45f76132017-02-22 00:06:47159 }
160 return ctx;
161}
162
163/*! POOL_join() :
164 Shutdown the queue, wake any sleeping threads, and join all of the threads.
165*/
ghatdevc3115f62018-01-25 16:01:31166static void POOL_join(POOL_ctx* ctx) {
Vianney Tran45f76132017-02-22 00:06:47167 /* Shut down the queue */
ghatdevc3115f62018-01-25 16:01:31168 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
Vianney Tran45f76132017-02-22 00:06:47169 ctx->shutdown = 1;
ghatdevc3115f62018-01-25 16:01:31170 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
Vianney Tran45f76132017-02-22 00:06:47171 /* Wake up sleeping threads */
ghatdevc3115f62018-01-25 16:01:31172 ZSTD_pthread_cond_broadcast(&ctx->queuePushCond);
173 ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
Vianney Tran45f76132017-02-22 00:06:47174 /* Join all of the threads */
175 { size_t i;
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34176 for (i = 0; i < ctx->threadCapacity; ++i) {
Vianney Tranca4d3c72023-04-20 20:30:32177 ZSTD_pthread_join(ctx->threads[i]); /* note : could fail */
Vianney Tran45f76132017-02-22 00:06:47178 } }
179}
180
181void POOL_free(POOL_ctx *ctx) {
182 if (!ctx) { return; }
183 POOL_join(ctx);
ghatdevc3115f62018-01-25 16:01:31184 ZSTD_pthread_mutex_destroy(&ctx->queueMutex);
185 ZSTD_pthread_cond_destroy(&ctx->queuePushCond);
186 ZSTD_pthread_cond_destroy(&ctx->queuePopCond);
Vianney Tran9f7d8ae2021-02-04 18:32:48187 ZSTD_customFree(ctx->queue, ctx->customMem);
188 ZSTD_customFree(ctx->threads, ctx->customMem);
189 ZSTD_customFree(ctx, ctx->customMem);
Vianney Tran45f76132017-02-22 00:06:47190}
191
Vianney Tranca4d3c72023-04-20 20:30:32192/*! POOL_joinJobs() :
193 * Waits for all queued jobs to finish executing.
194 */
195void POOL_joinJobs(POOL_ctx* ctx) {
196 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
197 while(!ctx->queueEmpty || ctx->numThreadsBusy > 0) {
198 ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex);
199 }
200 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
201}
202
Vianney Tran9f7d8ae2021-02-04 18:32:48203void ZSTD_freeThreadPool (ZSTD_threadPool* pool) {
204 POOL_free (pool);
205}
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34206
Vianney Tran7c9ad242022-04-14 16:08:30207size_t POOL_sizeof(const POOL_ctx* ctx) {
Elijah Andrews97f12e52017-08-10 19:14:23208 if (ctx==NULL) return 0; /* supports sizeof NULL */
209 return sizeof(*ctx)
210 + ctx->queueSize * sizeof(POOL_job)
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34211 + ctx->threadCapacity * sizeof(ZSTD_pthread_t);
212}
213
214
215/* @return : 0 on success, 1 on error */
216static int POOL_resize_internal(POOL_ctx* ctx, size_t numThreads)
217{
218 if (numThreads <= ctx->threadCapacity) {
219 if (!numThreads) return 1;
220 ctx->threadLimit = numThreads;
221 return 0;
222 }
223 /* numThreads > threadCapacity */
Vianney Tranca4d3c72023-04-20 20:30:32224 { ZSTD_pthread_t* const threadPool = (ZSTD_pthread_t*)ZSTD_customCalloc(numThreads * sizeof(ZSTD_pthread_t), ctx->customMem);
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34225 if (!threadPool) return 1;
226 /* replace existing thread pool */
Vianney Traned87d432024-07-26 20:24:58227 ZSTD_memcpy(threadPool, ctx->threads, ctx->threadCapacity * sizeof(ZSTD_pthread_t));
Vianney Tran9f7d8ae2021-02-04 18:32:48228 ZSTD_customFree(ctx->threads, ctx->customMem);
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34229 ctx->threads = threadPool;
230 /* Initialize additional threads */
231 { size_t threadId;
232 for (threadId = ctx->threadCapacity; threadId < numThreads; ++threadId) {
233 if (ZSTD_pthread_create(&threadPool[threadId], NULL, &POOL_thread, ctx)) {
234 ctx->threadCapacity = threadId;
235 return 1;
236 } }
237 } }
238 /* successfully expanded */
239 ctx->threadCapacity = numThreads;
240 ctx->threadLimit = numThreads;
241 return 0;
242}
243
244/* @return : 0 on success, 1 on error */
245int POOL_resize(POOL_ctx* ctx, size_t numThreads)
246{
247 int result;
248 if (ctx==NULL) return 1;
249 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
250 result = POOL_resize_internal(ctx, numThreads);
251 ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
252 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
253 return result;
Elijah Andrews97f12e52017-08-10 19:14:23254}
255
ghatdevc3115f62018-01-25 16:01:31256/**
257 * Returns 1 if the queue is full and 0 otherwise.
258 *
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34259 * When queueSize is 1 (pool was created with an intended queueSize of 0),
260 * then a queue is empty if there is a thread free _and_ no job is waiting.
ghatdevc3115f62018-01-25 16:01:31261 */
262static int isQueueFull(POOL_ctx const* ctx) {
263 if (ctx->queueSize > 1) {
264 return ctx->queueHead == ((ctx->queueTail + 1) % ctx->queueSize);
265 } else {
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34266 return (ctx->numThreadsBusy == ctx->threadLimit) ||
ghatdevc3115f62018-01-25 16:01:31267 !ctx->queueEmpty;
268 }
269}
270
Vianney Tran45f76132017-02-22 00:06:47271
Vianney Tran7c9ad242022-04-14 16:08:30272static void
273POOL_add_internal(POOL_ctx* ctx, POOL_function function, void *opaque)
Elijah Andrews80699a72018-03-30 19:22:31274{
Vianney Tranca4d3c72023-04-20 20:30:32275 POOL_job job;
276 job.function = function;
277 job.opaque = opaque;
Elijah Andrews80699a72018-03-30 19:22:31278 assert(ctx != NULL);
279 if (ctx->shutdown) return;
ghatdevc3115f62018-01-25 16:01:31280
Elijah Andrews80699a72018-03-30 19:22:31281 ctx->queueEmpty = 0;
282 ctx->queue[ctx->queueTail] = job;
283 ctx->queueTail = (ctx->queueTail + 1) % ctx->queueSize;
ghatdevc3115f62018-01-25 16:01:31284 ZSTD_pthread_cond_signal(&ctx->queuePopCond);
Vianney Tran45f76132017-02-22 00:06:47285}
286
Elijah Andrews80699a72018-03-30 19:22:31287void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque)
288{
289 assert(ctx != NULL);
290 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
291 /* Wait until there is space in the queue for the new job */
292 while (isQueueFull(ctx) && (!ctx->shutdown)) {
293 ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex);
294 }
295 POOL_add_internal(ctx, function, opaque);
296 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
297}
Vianney Tran45f76132017-02-22 00:06:47298
Elijah Andrews80699a72018-03-30 19:22:31299
300int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque)
301{
302 assert(ctx != NULL);
303 ZSTD_pthread_mutex_lock(&ctx->queueMutex);
304 if (isQueueFull(ctx)) {
305 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
306 return 0;
307 }
308 POOL_add_internal(ctx, function, opaque);
309 ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
310 return 1;
311}
312
313
314#else /* ZSTD_MULTITHREAD not defined */
315
316/* ========================== */
317/* No multi-threading support */
318/* ========================== */
319
320
321/* We don't need any data, but if it is empty, malloc() might return NULL. */
Vianney Tran45f76132017-02-22 00:06:47322struct POOL_ctx_s {
ghatdevc3115f62018-01-25 16:01:31323 int dummy;
Vianney Tran45f76132017-02-22 00:06:47324};
Vianney Tran9f7d8ae2021-02-04 18:32:48325static POOL_ctx g_poolCtx;
Vianney Tran45f76132017-02-22 00:06:47326
ghatdevc3115f62018-01-25 16:01:31327POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
328 return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
Vianney Tran45f76132017-02-22 00:06:47329}
330
Vianney Tran7c9ad242022-04-14 16:08:30331POOL_ctx*
332POOL_create_advanced(size_t numThreads, size_t queueSize, ZSTD_customMem customMem)
333{
ghatdevc3115f62018-01-25 16:01:31334 (void)numThreads;
335 (void)queueSize;
336 (void)customMem;
Vianney Tran9f7d8ae2021-02-04 18:32:48337 return &g_poolCtx;
Vianney Tran45f76132017-02-22 00:06:47338}
339
ghatdevc3115f62018-01-25 16:01:31340void POOL_free(POOL_ctx* ctx) {
Vianney Tran9f7d8ae2021-02-04 18:32:48341 assert(!ctx || ctx == &g_poolCtx);
ghatdevc3115f62018-01-25 16:01:31342 (void)ctx;
Vianney Tran45f76132017-02-22 00:06:47343}
344
Vianney Tranca4d3c72023-04-20 20:30:32345void POOL_joinJobs(POOL_ctx* ctx){
346 assert(!ctx || ctx == &g_poolCtx);
347 (void)ctx;
348}
349
Vlad Gorodetskyf38c0ed2019-04-08 11:45:34350int POOL_resize(POOL_ctx* ctx, size_t numThreads) {
351 (void)ctx; (void)numThreads;
352 return 0;
353}
354
Elijah Andrews80699a72018-03-30 19:22:31355void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) {
ghatdevc3115f62018-01-25 16:01:31356 (void)ctx;
357 function(opaque);
358}
359
Elijah Andrews80699a72018-03-30 19:22:31360int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque) {
361 (void)ctx;
362 function(opaque);
363 return 1;
364}
365
Vianney Tran7c9ad242022-04-14 16:08:30366size_t POOL_sizeof(const POOL_ctx* ctx) {
Elijah Andrews97f12e52017-08-10 19:14:23367 if (ctx==NULL) return 0; /* supports sizeof NULL */
Vianney Tran9f7d8ae2021-02-04 18:32:48368 assert(ctx == &g_poolCtx);
Elijah Andrews97f12e52017-08-10 19:14:23369 return sizeof(*ctx);
370}
371
Vianney Tran45f76132017-02-22 00:06:47372#endif /* ZSTD_MULTITHREAD */
Francois Visconte77df43e2021-11-18 12:49:40373
374#endif /* USE_EXTERNAL_ZSTD */