2
2
#define ZSTD_RUBY_H 1
3
3
4
4
#include <ruby.h>
5
+ #ifdef HAVE_RUBY_THREAD_H
6
+ #include <ruby/thread.h>
7
+ #endif
5
8
#include "./libzstd/zstd.h"
6
9
7
10
static int convert_compression_level (VALUE compression_level_value )
@@ -12,18 +15,40 @@ static int convert_compression_level(VALUE compression_level_value)
12
15
return NUM2INT (compression_level_value );
13
16
}
14
17
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
+
15
33
static size_t zstd_compress (ZSTD_CCtx * const ctx , ZSTD_outBuffer * output , ZSTD_inBuffer * input , ZSTD_EndDirective endOp )
16
34
{
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
18
42
}
19
43
20
44
static void set_compress_params (ZSTD_CCtx * const ctx , VALUE level_from_args , VALUE kwargs )
21
45
{
22
- ID kwargs_keys [2 ];
46
+ ID kwargs_keys [3 ];
23
47
kwargs_keys [0 ] = rb_intern ("level" );
24
48
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 );
27
52
28
53
int compression_level = ZSTD_CLEVEL_DEFAULT ;
29
54
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
43
68
rb_raise (rb_eRuntimeError , "%s" , "ZSTD_CCtx_loadDictionary failed" );
44
69
}
45
70
}
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
+ }
46
80
}
47
81
48
82
static void set_decompress_params (ZSTD_DCtx * const dctx , VALUE kwargs )
0 commit comments