libzstd_rs_sys/lib/common/
pool.rs

1use core::ptr;
2use std::sync::{Condvar, Mutex};
3use std::thread::JoinHandle;
4
5use libc::size_t;
6
7use crate::lib::common::allocations::{ZSTD_customCalloc, ZSTD_customFree};
8use crate::lib::zstd::ZSTD_customMem;
9
10pub struct POOL_ctx {
11    customMem: ZSTD_customMem,
12    threads: *mut JoinHandle<()>,
13    threadCapacity: size_t,
14    threadLimit: size_t,
15    queue: *mut POOL_job,
16    queueHead: size_t,
17    queueTail: size_t,
18    queueSize: size_t,
19    numThreadsBusy: size_t,
20    queueEmpty: core::ffi::c_int,
21    queueMutex: Mutex<()>,
22    queuePushCond: Condvar,
23    queuePopCond: Condvar,
24    shutdown: core::ffi::c_int,
25}
26
27struct SendPoolCtxPtr(*mut POOL_ctx);
28unsafe impl Send for SendPoolCtxPtr {}
29
30#[derive(Copy, Clone)]
31#[repr(C)]
32pub(crate) struct POOL_job {
33    function: POOL_function,
34    opaque: *mut core::ffi::c_void,
35}
36pub type POOL_function = unsafe fn(*mut core::ffi::c_void) -> ();
37pub type ZSTD_threadPool = POOL_ctx;
38
39unsafe fn POOL_thread(ctx: *mut POOL_ctx) {
40    if ctx.is_null() {
41        return;
42    }
43    loop {
44        let mut guard = (*ctx).queueMutex.lock().unwrap();
45        #[expect(clippy::while_immutable_condition)]
46        while (*ctx).queueEmpty != 0 || (*ctx).numThreadsBusy >= (*ctx).threadLimit {
47            if (*ctx).shutdown != 0 {
48                return;
49            }
50            guard = (*ctx).queuePopCond.wait(guard).unwrap();
51        }
52        let job = *((*ctx).queue).add((*ctx).queueHead);
53        (*ctx).queueHead = ((*ctx).queueHead).wrapping_add(1) % (*ctx).queueSize;
54        (*ctx).numThreadsBusy += 1;
55        (*ctx).queueEmpty = ((*ctx).queueHead == (*ctx).queueTail) as core::ffi::c_int;
56        (*ctx).queuePushCond.notify_one();
57        drop(guard);
58        (job.function)(job.opaque);
59        guard = (*ctx).queueMutex.lock().unwrap();
60        (*ctx).numThreadsBusy -= 1;
61        (*ctx).queuePushCond.notify_one();
62    }
63}
64#[cfg_attr(feature = "export-symbols", export_name = crate::prefix!(ZSTD_createThreadPool))]
65pub unsafe extern "C" fn ZSTD_createThreadPool(numThreads: size_t) -> *mut ZSTD_threadPool {
66    POOL_create(numThreads, 0)
67}
68pub unsafe fn POOL_create(numThreads: size_t, queueSize: size_t) -> *mut POOL_ctx {
69    POOL_create_advanced(numThreads, queueSize, ZSTD_customMem::default())
70}
71pub(crate) unsafe fn POOL_create_advanced(
72    numThreads: size_t,
73    queueSize: size_t,
74    customMem: ZSTD_customMem,
75) -> *mut POOL_ctx {
76    let mut ctx = core::ptr::null_mut::<POOL_ctx>();
77    if numThreads == 0 {
78        return core::ptr::null_mut();
79    }
80    ctx = ZSTD_customCalloc(::core::mem::size_of::<POOL_ctx>(), customMem) as *mut POOL_ctx;
81    if ctx.is_null() {
82        return core::ptr::null_mut();
83    }
84    (*ctx).queueSize = queueSize.wrapping_add(1);
85    (*ctx).queue = ZSTD_customCalloc(
86        (*ctx).queueSize * ::core::mem::size_of::<POOL_job>(),
87        customMem,
88    ) as *mut POOL_job;
89    (*ctx).queueHead = 0;
90    (*ctx).queueTail = 0;
91    (*ctx).numThreadsBusy = 0;
92    (*ctx).queueEmpty = 1;
93    ptr::write(ptr::addr_of_mut!((*ctx).queueMutex), Mutex::new(()));
94    ptr::write(ptr::addr_of_mut!((*ctx).queuePushCond), Condvar::new());
95    ptr::write(ptr::addr_of_mut!((*ctx).queuePopCond), Condvar::new());
96    (*ctx).shutdown = 0;
97    (*ctx).threads = ZSTD_customCalloc(
98        numThreads * ::core::mem::size_of::<JoinHandle<()>>(),
99        customMem,
100    ) as *mut JoinHandle<()>;
101    (*ctx).threadCapacity = 0;
102    (*ctx).customMem = customMem;
103    if (*ctx).threads.is_null() || (*ctx).queue.is_null() {
104        POOL_free(ctx);
105        return core::ptr::null_mut();
106    }
107    for i in 0..numThreads {
108        let ctx = SendPoolCtxPtr(ctx);
109        core::ptr::write(
110            (*ctx.0).threads.add(i),
111            std::thread::spawn(|| {
112                let ctx = ctx;
113                POOL_thread(ctx.0)
114            }),
115        );
116    }
117    (*ctx).threadCapacity = numThreads;
118    (*ctx).threadLimit = numThreads;
119    ctx
120}
121unsafe fn POOL_join(ctx: *mut POOL_ctx) {
122    let guard = (*ctx).queueMutex.lock().unwrap();
123    (*ctx).shutdown = 1;
124    drop(guard);
125    (*ctx).queuePushCond.notify_all();
126    (*ctx).queuePopCond.notify_all();
127    for i in 0..(*ctx).threadCapacity {
128        core::ptr::read((*ctx).threads.add(i)).join().unwrap();
129    }
130}
131pub unsafe fn POOL_free(ctx: *mut POOL_ctx) {
132    if ctx.is_null() {
133        return;
134    }
135    POOL_join(ctx);
136    ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queueMutex));
137    ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queuePushCond));
138    ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queuePopCond));
139    ZSTD_customFree(
140        (*ctx).queue as *mut core::ffi::c_void,
141        (*ctx).queueSize * ::core::mem::size_of::<POOL_job>(),
142        (*ctx).customMem,
143    );
144    ZSTD_customFree(
145        (*ctx).threads as *mut core::ffi::c_void,
146        (*ctx).threadCapacity * ::core::mem::size_of::<JoinHandle<()>>(),
147        (*ctx).customMem,
148    );
149    ZSTD_customFree(
150        ctx as *mut core::ffi::c_void,
151        ::core::mem::size_of::<POOL_ctx>(),
152        (*ctx).customMem,
153    );
154}
155pub unsafe fn POOL_joinJobs(ctx: *mut POOL_ctx) {
156    let mut guard = (*ctx).queueMutex.lock().unwrap();
157    #[expect(clippy::while_immutable_condition)]
158    while (*ctx).queueEmpty == 0 || (*ctx).numThreadsBusy > 0 {
159        guard = (*ctx).queuePushCond.wait(guard).unwrap();
160    }
161}
162#[cfg_attr(feature = "export-symbols", export_name = crate::prefix!(ZSTD_freeThreadPool))]
163pub unsafe extern "C" fn ZSTD_freeThreadPool(pool: *mut ZSTD_threadPool) {
164    POOL_free(pool);
165}
166pub(crate) unsafe fn POOL_sizeof(ctx: *const POOL_ctx) -> size_t {
167    if ctx.is_null() {
168        return 0;
169    }
170    ::core::mem::size_of::<POOL_ctx>()
171        + (*ctx).queueSize * ::core::mem::size_of::<POOL_job>()
172        + (*ctx).threadCapacity * ::core::mem::size_of::<JoinHandle<()>>()
173}
174unsafe fn POOL_resize_internal(ctx: *mut POOL_ctx, numThreads: size_t) -> core::ffi::c_int {
175    if numThreads <= (*ctx).threadCapacity {
176        if numThreads == 0 {
177            return 1;
178        }
179        (*ctx).threadLimit = numThreads;
180        return 0;
181    }
182    let threadPool = ZSTD_customCalloc(
183        numThreads.wrapping_mul(::core::mem::size_of::<JoinHandle<()>>()),
184        (*ctx).customMem,
185    ) as *mut JoinHandle<()>;
186    if threadPool.is_null() {
187        return 1;
188    }
189    libc::memcpy(
190        threadPool as *mut core::ffi::c_void,
191        (*ctx).threads as *const core::ffi::c_void,
192        (*ctx).threadCapacity * ::core::mem::size_of::<JoinHandle<()>>(),
193    );
194    ZSTD_customFree(
195        (*ctx).threads as *mut core::ffi::c_void,
196        (*ctx).threadCapacity * ::core::mem::size_of::<JoinHandle<()>>(),
197        (*ctx).customMem,
198    );
199    (*ctx).threads = threadPool;
200    for threadId in (*ctx).threadCapacity..numThreads {
201        let ctx = SendPoolCtxPtr(ctx);
202        core::ptr::write(
203            ((*ctx.0).threads).add(threadId),
204            std::thread::spawn(|| {
205                let ctx = ctx;
206                POOL_thread(ctx.0)
207            }),
208        );
209    }
210    (*ctx).threadCapacity = numThreads;
211    (*ctx).threadLimit = numThreads;
212    0
213}
214pub(crate) unsafe fn POOL_resize(ctx: *mut POOL_ctx, numThreads: size_t) -> core::ffi::c_int {
215    if ctx.is_null() {
216        return 1;
217    }
218    let _guard = (*ctx).queueMutex.lock().unwrap();
219    let result = POOL_resize_internal(ctx, numThreads);
220    (*ctx).queuePopCond.notify_all();
221    result
222}
223unsafe fn isQueueFull(ctx: *const POOL_ctx) -> core::ffi::c_int {
224    if (*ctx).queueSize > 1 {
225        ((*ctx).queueHead == ((*ctx).queueTail).wrapping_add(1) % (*ctx).queueSize)
226            as core::ffi::c_int
227    } else {
228        ((*ctx).numThreadsBusy == (*ctx).threadLimit || (*ctx).queueEmpty == 0) as core::ffi::c_int
229    }
230}
231unsafe fn POOL_add_internal(
232    ctx: *mut POOL_ctx,
233    function: POOL_function,
234    opaque: *mut core::ffi::c_void,
235) {
236    let job = POOL_job { function, opaque };
237    if (*ctx).shutdown != 0 {
238        return;
239    }
240    (*ctx).queueEmpty = 0;
241    *((*ctx).queue).add((*ctx).queueTail) = job;
242    (*ctx).queueTail = ((*ctx).queueTail).wrapping_add(1) % (*ctx).queueSize;
243    (*ctx).queuePopCond.notify_one();
244}
245pub unsafe fn POOL_add(
246    ctx: *mut POOL_ctx,
247    function: POOL_function,
248    opaque: *mut core::ffi::c_void,
249) {
250    let mut guard = (*ctx).queueMutex.lock().unwrap();
251    while isQueueFull(ctx) != 0 && (*ctx).shutdown == 0 {
252        guard = (*ctx).queuePushCond.wait(guard).unwrap();
253    }
254    POOL_add_internal(ctx, function, opaque);
255}
256pub(crate) unsafe fn POOL_tryAdd(
257    ctx: *mut POOL_ctx,
258    function: POOL_function,
259    opaque: *mut core::ffi::c_void,
260) -> core::ffi::c_int {
261    let _guard = (*ctx).queueMutex.lock().unwrap();
262    if isQueueFull(ctx) != 0 {
263        return 0;
264    }
265    POOL_add_internal(ctx, function, opaque);
266    1
267}
libzstd_rs_sys/lib/common/pool.rs

libzstd_rs_sys/lib/common/
pool.rs
