libzstd_rs_sys/lib/common/
pool.rs

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

libzstd_rs_sys/lib/common/
pool.rs
