pool.c 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341
  1. /*
  2. * Copyright (c) 2016-present, Yann Collet, Facebook, Inc.
  3. * All rights reserved.
  4. *
  5. * This source code is licensed under both the BSD-style license (found in the
  6. * LICENSE file in the root directory of this source tree) and the GPLv2 (found
  7. * in the COPYING file in the root directory of this source tree).
  8. * You may select, at your option, one of the above-listed licenses.
  9. */
  10. /* ====== Dependencies ======= */
  11. #include <stddef.h> /* size_t */
  12. #include "debug.h" /* assert */
  13. #include "zstd_internal.h" /* ZSTD_malloc, ZSTD_free */
  14. #include "pool.h"
  15. /* ====== Compiler specifics ====== */
  16. #if defined(_MSC_VER)
  17. # pragma warning(disable : 4204) /* disable: C4204: non-constant aggregate initializer */
  18. #endif
  19. #ifdef ZSTD_MULTITHREAD
  20. #include "threading.h" /* pthread adaptation */
  21. /* A job is a function and an opaque argument */
  22. typedef struct POOL_job_s {
  23. POOL_function function;
  24. void *opaque;
  25. } POOL_job;
  26. struct POOL_ctx_s {
  27. ZSTD_customMem customMem;
  28. /* Keep track of the threads */
  29. ZSTD_pthread_t* threads;
  30. size_t threadCapacity;
  31. size_t threadLimit;
  32. /* The queue is a circular buffer */
  33. POOL_job *queue;
  34. size_t queueHead;
  35. size_t queueTail;
  36. size_t queueSize;
  37. /* The number of threads working on jobs */
  38. size_t numThreadsBusy;
  39. /* Indicates if the queue is empty */
  40. int queueEmpty;
  41. /* The mutex protects the queue */
  42. ZSTD_pthread_mutex_t queueMutex;
  43. /* Condition variable for pushers to wait on when the queue is full */
  44. ZSTD_pthread_cond_t queuePushCond;
  45. /* Condition variables for poppers to wait on when the queue is empty */
  46. ZSTD_pthread_cond_t queuePopCond;
  47. /* Indicates if the queue is shutting down */
  48. int shutdown;
  49. };
  50. /* POOL_thread() :
  51. * Work thread for the thread pool.
  52. * Waits for jobs and executes them.
  53. * @returns : NULL on failure else non-null.
  54. */
  55. static void* POOL_thread(void* opaque) {
  56. POOL_ctx* const ctx = (POOL_ctx*)opaque;
  57. if (!ctx) { return NULL; }
  58. for (;;) {
  59. /* Lock the mutex and wait for a non-empty queue or until shutdown */
  60. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  61. while ( ctx->queueEmpty
  62. || (ctx->numThreadsBusy >= ctx->threadLimit) ) {
  63. if (ctx->shutdown) {
  64. /* even if !queueEmpty, (possible if numThreadsBusy >= threadLimit),
  65. * a few threads will be shutdown while !queueEmpty,
  66. * but enough threads will remain active to finish the queue */
  67. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  68. return opaque;
  69. }
  70. ZSTD_pthread_cond_wait(&ctx->queuePopCond, &ctx->queueMutex);
  71. }
  72. /* Pop a job off the queue */
  73. { POOL_job const job = ctx->queue[ctx->queueHead];
  74. ctx->queueHead = (ctx->queueHead + 1) % ctx->queueSize;
  75. ctx->numThreadsBusy++;
  76. ctx->queueEmpty = ctx->queueHead == ctx->queueTail;
  77. /* Unlock the mutex, signal a pusher, and run the job */
  78. ZSTD_pthread_cond_signal(&ctx->queuePushCond);
  79. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  80. job.function(job.opaque);
  81. /* If the intended queue size was 0, signal after finishing job */
  82. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  83. ctx->numThreadsBusy--;
  84. if (ctx->queueSize == 1) {
  85. ZSTD_pthread_cond_signal(&ctx->queuePushCond);
  86. }
  87. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  88. }
  89. } /* for (;;) */
  90. assert(0); /* Unreachable */
  91. }
  92. POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
  93. return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
  94. }
  95. POOL_ctx* POOL_create_advanced(size_t numThreads, size_t queueSize,
  96. ZSTD_customMem customMem) {
  97. POOL_ctx* ctx;
  98. /* Check parameters */
  99. if (!numThreads) { return NULL; }
  100. /* Allocate the context and zero initialize */
  101. ctx = (POOL_ctx*)ZSTD_calloc(sizeof(POOL_ctx), customMem);
  102. if (!ctx) { return NULL; }
  103. /* Initialize the job queue.
  104. * It needs one extra space since one space is wasted to differentiate
  105. * empty and full queues.
  106. */
  107. ctx->queueSize = queueSize + 1;
  108. ctx->queue = (POOL_job*)ZSTD_malloc(ctx->queueSize * sizeof(POOL_job), customMem);
  109. ctx->queueHead = 0;
  110. ctx->queueTail = 0;
  111. ctx->numThreadsBusy = 0;
  112. ctx->queueEmpty = 1;
  113. (void)ZSTD_pthread_mutex_init(&ctx->queueMutex, NULL);
  114. (void)ZSTD_pthread_cond_init(&ctx->queuePushCond, NULL);
  115. (void)ZSTD_pthread_cond_init(&ctx->queuePopCond, NULL);
  116. ctx->shutdown = 0;
  117. /* Allocate space for the thread handles */
  118. ctx->threads = (ZSTD_pthread_t*)ZSTD_malloc(numThreads * sizeof(ZSTD_pthread_t), customMem);
  119. ctx->threadCapacity = 0;
  120. ctx->customMem = customMem;
  121. /* Check for errors */
  122. if (!ctx->threads || !ctx->queue) { POOL_free(ctx); return NULL; }
  123. /* Initialize the threads */
  124. { size_t i;
  125. for (i = 0; i < numThreads; ++i) {
  126. if (ZSTD_pthread_create(&ctx->threads[i], NULL, &POOL_thread, ctx)) {
  127. ctx->threadCapacity = i;
  128. POOL_free(ctx);
  129. return NULL;
  130. } }
  131. ctx->threadCapacity = numThreads;
  132. ctx->threadLimit = numThreads;
  133. }
  134. return ctx;
  135. }
  136. /*! POOL_join() :
  137. Shutdown the queue, wake any sleeping threads, and join all of the threads.
  138. */
  139. static void POOL_join(POOL_ctx* ctx) {
  140. /* Shut down the queue */
  141. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  142. ctx->shutdown = 1;
  143. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  144. /* Wake up sleeping threads */
  145. ZSTD_pthread_cond_broadcast(&ctx->queuePushCond);
  146. ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
  147. /* Join all of the threads */
  148. { size_t i;
  149. for (i = 0; i < ctx->threadCapacity; ++i) {
  150. ZSTD_pthread_join(ctx->threads[i], NULL); /* note : could fail */
  151. } }
  152. }
  153. void POOL_free(POOL_ctx *ctx) {
  154. if (!ctx) { return; }
  155. POOL_join(ctx);
  156. ZSTD_pthread_mutex_destroy(&ctx->queueMutex);
  157. ZSTD_pthread_cond_destroy(&ctx->queuePushCond);
  158. ZSTD_pthread_cond_destroy(&ctx->queuePopCond);
  159. ZSTD_free(ctx->queue, ctx->customMem);
  160. ZSTD_free(ctx->threads, ctx->customMem);
  161. ZSTD_free(ctx, ctx->customMem);
  162. }
  163. size_t POOL_sizeof(POOL_ctx *ctx) {
  164. if (ctx==NULL) return 0; /* supports sizeof NULL */
  165. return sizeof(*ctx)
  166. + ctx->queueSize * sizeof(POOL_job)
  167. + ctx->threadCapacity * sizeof(ZSTD_pthread_t);
  168. }
  169. /* @return : 0 on success, 1 on error */
  170. static int POOL_resize_internal(POOL_ctx* ctx, size_t numThreads)
  171. {
  172. if (numThreads <= ctx->threadCapacity) {
  173. if (!numThreads) return 1;
  174. ctx->threadLimit = numThreads;
  175. return 0;
  176. }
  177. /* numThreads > threadCapacity */
  178. { ZSTD_pthread_t* const threadPool = (ZSTD_pthread_t*)ZSTD_malloc(numThreads * sizeof(ZSTD_pthread_t), ctx->customMem);
  179. if (!threadPool) return 1;
  180. /* replace existing thread pool */
  181. memcpy(threadPool, ctx->threads, ctx->threadCapacity * sizeof(*threadPool));
  182. ZSTD_free(ctx->threads, ctx->customMem);
  183. ctx->threads = threadPool;
  184. /* Initialize additional threads */
  185. { size_t threadId;
  186. for (threadId = ctx->threadCapacity; threadId < numThreads; ++threadId) {
  187. if (ZSTD_pthread_create(&threadPool[threadId], NULL, &POOL_thread, ctx)) {
  188. ctx->threadCapacity = threadId;
  189. return 1;
  190. } }
  191. } }
  192. /* successfully expanded */
  193. ctx->threadCapacity = numThreads;
  194. ctx->threadLimit = numThreads;
  195. return 0;
  196. }
  197. /* @return : 0 on success, 1 on error */
  198. int POOL_resize(POOL_ctx* ctx, size_t numThreads)
  199. {
  200. int result;
  201. if (ctx==NULL) return 1;
  202. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  203. result = POOL_resize_internal(ctx, numThreads);
  204. ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
  205. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  206. return result;
  207. }
  208. /**
  209. * Returns 1 if the queue is full and 0 otherwise.
  210. *
  211. * When queueSize is 1 (pool was created with an intended queueSize of 0),
  212. * then a queue is empty if there is a thread free _and_ no job is waiting.
  213. */
  214. static int isQueueFull(POOL_ctx const* ctx) {
  215. if (ctx->queueSize > 1) {
  216. return ctx->queueHead == ((ctx->queueTail + 1) % ctx->queueSize);
  217. } else {
  218. return (ctx->numThreadsBusy == ctx->threadLimit) ||
  219. !ctx->queueEmpty;
  220. }
  221. }
  222. static void POOL_add_internal(POOL_ctx* ctx, POOL_function function, void *opaque)
  223. {
  224. POOL_job const job = {function, opaque};
  225. assert(ctx != NULL);
  226. if (ctx->shutdown) return;
  227. ctx->queueEmpty = 0;
  228. ctx->queue[ctx->queueTail] = job;
  229. ctx->queueTail = (ctx->queueTail + 1) % ctx->queueSize;
  230. ZSTD_pthread_cond_signal(&ctx->queuePopCond);
  231. }
  232. void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque)
  233. {
  234. assert(ctx != NULL);
  235. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  236. /* Wait until there is space in the queue for the new job */
  237. while (isQueueFull(ctx) && (!ctx->shutdown)) {
  238. ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex);
  239. }
  240. POOL_add_internal(ctx, function, opaque);
  241. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  242. }
  243. int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque)
  244. {
  245. assert(ctx != NULL);
  246. ZSTD_pthread_mutex_lock(&ctx->queueMutex);
  247. if (isQueueFull(ctx)) {
  248. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  249. return 0;
  250. }
  251. POOL_add_internal(ctx, function, opaque);
  252. ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
  253. return 1;
  254. }
  255. #else /* ZSTD_MULTITHREAD not defined */
  256. /* ========================== */
  257. /* No multi-threading support */
  258. /* ========================== */
  259. /* We don't need any data, but if it is empty, malloc() might return NULL. */
  260. struct POOL_ctx_s {
  261. int dummy;
  262. };
  263. static POOL_ctx g_ctx;
  264. POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
  265. return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
  266. }
  267. POOL_ctx* POOL_create_advanced(size_t numThreads, size_t queueSize, ZSTD_customMem customMem) {
  268. (void)numThreads;
  269. (void)queueSize;
  270. (void)customMem;
  271. return &g_ctx;
  272. }
  273. void POOL_free(POOL_ctx* ctx) {
  274. assert(!ctx || ctx == &g_ctx);
  275. (void)ctx;
  276. }
  277. int POOL_resize(POOL_ctx* ctx, size_t numThreads) {
  278. (void)ctx; (void)numThreads;
  279. return 0;
  280. }
  281. void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) {
  282. (void)ctx;
  283. function(opaque);
  284. }
  285. int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque) {
  286. (void)ctx;
  287. function(opaque);
  288. return 1;
  289. }
  290. size_t POOL_sizeof(POOL_ctx* ctx) {
  291. if (ctx==NULL) return 0; /* supports sizeof NULL */
  292. assert(ctx == &g_ctx);
  293. return sizeof(*ctx);
  294. }
  295. #endif /* ZSTD_MULTITHREAD */