22#define ZSTD_RUBY_H 1
33
44#include <ruby.h>
5+ #ifdef HAVE_RUBY_THREAD_H
6+ #include <ruby/thread.h>
7+ #endif
58#include "./libzstd/zstd.h"
69
710static int convert_compression_level (VALUE compression_level_value )
@@ -12,18 +15,40 @@ static int convert_compression_level(VALUE compression_level_value)
1215 return NUM2INT (compression_level_value );
1316}
1417
18+ struct compress_params {
19+ ZSTD_CCtx * ctx ;
20+ ZSTD_outBuffer * output ;
21+ ZSTD_inBuffer * input ;
22+ ZSTD_EndDirective endOp ;
23+ size_t ret ;
24+ };
25+
26+ static void * compress_wrapper (void * args )
27+ {
28+ struct compress_params * params = args ;
29+ params -> ret = ZSTD_compressStream2 (params -> ctx , params -> output , params -> input , params -> endOp );
30+ return NULL ;
31+ }
32+
1533static size_t zstd_compress (ZSTD_CCtx * const ctx , ZSTD_outBuffer * output , ZSTD_inBuffer * input , ZSTD_EndDirective endOp )
1634{
17- return ZSTD_compressStream2 (ctx , output , input , endOp );
35+ #ifdef HAVE_RUBY_THREAD_H
36+ struct compress_params params = { ctx , output , input , endOp };
37+ rb_thread_call_without_gvl (compress_wrapper , & params , RUBY_UBF_IO , NULL );
38+ return params .ret ;
39+ #else
40+ return ZSTD_compressStream2 (ctx , output , input , endOp );
41+ #endif
1842}
1943
2044static void set_compress_params (ZSTD_CCtx * const ctx , VALUE level_from_args , VALUE kwargs )
2145{
22- ID kwargs_keys [2 ];
46+ ID kwargs_keys [3 ];
2347 kwargs_keys [0 ] = rb_intern ("level" );
2448 kwargs_keys [1 ] = rb_intern ("dict" );
25- VALUE kwargs_values [2 ];
26- rb_get_kwargs (kwargs , kwargs_keys , 0 , 2 , kwargs_values );
49+ kwargs_keys [2 ] = rb_intern ("thread_num" );
50+ VALUE kwargs_values [3 ];
51+ rb_get_kwargs (kwargs , kwargs_keys , 0 , 3 , kwargs_values );
2752
2853 int compression_level = ZSTD_CLEVEL_DEFAULT ;
2954 if (kwargs_values [0 ] != Qundef && kwargs_values [0 ] != Qnil ) {
@@ -43,6 +68,15 @@ static void set_compress_params(ZSTD_CCtx* const ctx, VALUE level_from_args, VAL
4368 rb_raise (rb_eRuntimeError , "%s" , "ZSTD_CCtx_loadDictionary failed" );
4469 }
4570 }
71+
72+ if (kwargs_values [2 ] != Qundef && kwargs_values [2 ] != Qnil ) {
73+ int thread_num = NUM2INT (kwargs_values [2 ]);
74+ size_t const r = ZSTD_CCtx_setParameter (ctx , ZSTD_c_nbWorkers , thread_num );
75+ if (ZSTD_isError (r )) {
76+ rb_warn ("Note: the linked libzstd library doesn't support multithreading.Reverting to single-thread mode. \n" );
77+ }
78+ // ZSTD_CCtx_setParameter(ctx, ZSTD_c_jobSize, thread_num);
79+ }
4680}
4781
4882static void set_decompress_params (ZSTD_DCtx * const dctx , VALUE kwargs )
0 commit comments