Merge pull request #53 from tidewise/nogvl

SpringMT · web-flow · commit 4df4d623e08a · 2023-06-07T00:25:09.000+09:00
unlock the GVL for compression/decompression operations
diff --git a/ext/zstdruby/streaming_compress.c b/ext/zstdruby/streaming_compress.c
@@ -1,10 +1,12 @@
 #include <common.h>
 #include <streaming_compress.h>
+#include <ruby/thread.h>
 
 struct streaming_compress_t {
   ZSTD_CCtx* ctx;
   VALUE buf;
   size_t buf_size;
+  char nogvl;
 };
 
 static void
@@ -52,11 +54,18 @@ static VALUE
 rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
 {
   VALUE compression_level_value;
-  rb_scan_args(argc, argv, "01", &compression_level_value);
+  VALUE kwargs;
+  rb_scan_args(argc, argv, "01:", &compression_level_value, &kwargs);
   int compression_level = convert_compression_level(compression_level_value);
 
+  ID kwargs_keys[1];
+  kwargs_keys[0] = rb_intern("no_gvl");
+  VALUE kwargs_values[1];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+
   struct streaming_compress_t* sc;
   TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  sc->nogvl = kwargs_values[0] != Qundef && RTEST(kwargs_values[0]);
   size_t const buffOutSize = ZSTD_CStreamOutSize();
 
   ZSTD_CCtx* ctx = ZSTD_createCCtx();
@@ -76,6 +85,35 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
     : (FIX2INT((val))))
 #define ARG_CONTINUE(val)     FIXNUMARG((val), ZSTD_e_continue)
 
+struct compress_stream_nogvl_t {
+  ZSTD_CCtx* ctx;
+  ZSTD_outBuffer* output;
+  ZSTD_inBuffer* input;
+  ZSTD_EndDirective endOp;
+  size_t ret;
+};
+
+static void*
+compressStream2_nogvl(void* arg)
+{
+  struct compress_stream_nogvl_t* params = arg;
+  params->ret = ZSTD_compressStream2(params->ctx, params->output, params->input, params->endOp);
+  return NULL;
+}
+
+static size_t
+compressStream2(char nogvl, ZSTD_CCtx* ctx, ZSTD_outBuffer* output, ZSTD_inBuffer* input, ZSTD_EndDirective endOp)
+{
+  struct compress_stream_nogvl_t params = { ctx, output, input, endOp, 0 };
+  if (nogvl) {
+    rb_thread_call_without_gvl(compressStream2_nogvl, &params, NULL, NULL);
+  }
+  else {
+    compressStream2_nogvl(&params);
+  }
+  return params.ret;
+}
+
 static VALUE
 no_compress(struct streaming_compress_t* sc, ZSTD_EndDirective endOp)
 {
@@ -86,7 +124,7 @@ no_compress(struct streaming_compress_t* sc, ZSTD_EndDirective endOp)
   do {
     ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
 
-    size_t const ret = ZSTD_compressStream2(sc->ctx, &output, &input, endOp);
+    size_t const ret = compressStream2(sc->nogvl, sc->ctx, &output, &input, endOp);
     if (ZSTD_isError(ret)) {
       rb_raise(rb_eRuntimeError, "flush error error code: %s", ZSTD_getErrorName(ret));
     }
@@ -109,7 +147,7 @@ rb_streaming_compress_compress(VALUE obj, VALUE src)
   VALUE result = rb_str_new(0, 0);
   while (input.pos < input.size) {
     ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
-    size_t const ret = ZSTD_compressStream2(sc->ctx, &output, &input, ZSTD_e_continue);
+    size_t const ret = compressStream2(sc->nogvl, sc->ctx, &output, &input, ZSTD_e_continue);
     if (ZSTD_isError(ret)) {
       rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret));
     }
@@ -132,7 +170,7 @@ rb_streaming_compress_addstr(VALUE obj, VALUE src)
 
   while (input.pos < input.size) {
     ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
-    size_t const result = ZSTD_compressStream2(sc->ctx, &output, &input, ZSTD_e_continue);
+    size_t const result = compressStream2(sc->nogvl, sc->ctx, &output, &input, ZSTD_e_continue);
     if (ZSTD_isError(result)) {
       rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(result));
     }
diff --git a/ext/zstdruby/streaming_decompress.c b/ext/zstdruby/streaming_decompress.c
@@ -1,9 +1,11 @@
 #include <common.h>
+#include <ruby/thread.h>
 
 struct streaming_decompress_t {
   ZSTD_DCtx* ctx;
   VALUE buf;
   size_t buf_size;
+  char nogvl;
 };
 
 static void
@@ -48,10 +50,19 @@ rb_streaming_decompress_allocate(VALUE klass)
 }
 
 static VALUE
-rb_streaming_decompress_initialize(VALUE obj)
+rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj)
 {
+  VALUE kwargs;
+  rb_scan_args(argc, argv, "00:", &kwargs);
+
+  ID kwargs_keys[1];
+  kwargs_keys[0] = rb_intern("no_gvl");
+  VALUE kwargs_values[1];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+
   struct streaming_decompress_t* sd;
   TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  sd->nogvl = kwargs_values[0] != Qundef && RTEST(kwargs_values[0]);
   size_t const buffOutSize = ZSTD_DStreamOutSize();
 
   ZSTD_DCtx* ctx = ZSTD_createDCtx();
@@ -65,6 +76,34 @@ rb_streaming_decompress_initialize(VALUE obj)
   return obj;
 }
 
+struct decompress_stream_nogvl_t {
+  ZSTD_DCtx* ctx;
+  ZSTD_outBuffer* output;
+  ZSTD_inBuffer* input;
+  size_t ret;
+};
+
+static void*
+decompressStream_nogvl(void* args)
+{
+  struct decompress_stream_nogvl_t* params = args;
+  params->ret = ZSTD_decompressStream(params->ctx, params->output, params->input);
+  return NULL;
+}
+
+static size_t
+decompressStream(char nogvl, ZSTD_DCtx* ctx, ZSTD_outBuffer* output, ZSTD_inBuffer* input)
+{
+  struct decompress_stream_nogvl_t params = { ctx, output, input, 0 };
+  if (nogvl) {
+    rb_thread_call_without_gvl(decompressStream_nogvl, &params, NULL, NULL);
+  }
+  else {
+    decompressStream_nogvl(&params);
+  }
+  return params.ret;
+}
+
 static VALUE
 rb_streaming_decompress_decompress(VALUE obj, VALUE src)
 {
@@ -79,7 +118,7 @@ rb_streaming_decompress_decompress(VALUE obj, VALUE src)
   VALUE result = rb_str_new(0, 0);
   while (input.pos < input.size) {
     ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
-    size_t const ret = ZSTD_decompressStream(sd->ctx, &output, &input);
+    size_t const ret = decompressStream(sd->nogvl, sd->ctx, &output, &input);
     if (ZSTD_isError(ret)) {
       rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret));
     }
@@ -102,7 +141,7 @@ rb_streaming_decompress_addstr(VALUE obj, VALUE src)
 
   while (input.pos < input.size) {
     ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
-    size_t const result = ZSTD_decompressStream(sd->ctx, &output, &input);
+    size_t const result = decompressStream(sd->nogvl, sd->ctx, &output, &input);
     if (ZSTD_isError(result)) {
       rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(result));
     }
@@ -116,7 +155,7 @@ zstd_ruby_streaming_decompress_init(void)
 {
   VALUE cStreamingDecompress = rb_define_class_under(rb_mZstd, "StreamingDecompress", rb_cObject);
   rb_define_alloc_func(cStreamingDecompress, rb_streaming_decompress_allocate);
-  rb_define_method(cStreamingDecompress, "initialize", rb_streaming_decompress_initialize, 0);
+  rb_define_method(cStreamingDecompress, "initialize", rb_streaming_decompress_initialize, -1);
   rb_define_method(cStreamingDecompress, "decompress", rb_streaming_decompress_decompress, 1);
   rb_define_method(cStreamingDecompress, "<<", rb_streaming_decompress_addstr, 1);
 }
diff --git a/spec/zstd-ruby-streaming-compress_spec.rb b/spec/zstd-ruby-streaming-compress_spec.rb
@@ -1,10 +1,10 @@
 require "spec_helper"
 require 'zstd-ruby'
 
-RSpec.describe Zstd::StreamingCompress do
+shared_examples "a streaming compressor" do
   describe '<<' do
     it 'shoud work' do
-      stream = Zstd::StreamingCompress.new
+      stream = Zstd::StreamingCompress.new(no_gvl: no_gvl)
       stream << "abc" << "def"
       res = stream.finish
       expect(Zstd.decompress(res)).to eq('abcdef')
@@ -13,7 +13,7 @@
 
   describe '<< + GC.compat' do
     it 'shoud work' do
-      stream = Zstd::StreamingCompress.new
+      stream = Zstd::StreamingCompress.new(no_gvl: no_gvl)
       stream << "abc" << "def"
       GC.compact
       stream << "ghi"
@@ -24,7 +24,7 @@
 
   describe '<< + flush' do
     it 'shoud work' do
-      stream = Zstd::StreamingCompress.new
+      stream = Zstd::StreamingCompress.new(no_gvl: no_gvl)
       stream << "abc" << "def"
       res = stream.flush
       stream << "ghi"
@@ -35,7 +35,7 @@
 
   describe 'compress + flush' do
     it 'shoud work' do
-      stream = Zstd::StreamingCompress.new
+      stream = Zstd::StreamingCompress.new(no_gvl: no_gvl)
       res = stream.compress("abc")
       res << stream.flush
       res << stream.compress("def")
@@ -46,7 +46,7 @@
 
   describe 'compression level' do
     it 'shoud work' do
-      stream = Zstd::StreamingCompress.new(5)
+      stream = Zstd::StreamingCompress.new(5, no_gvl: no_gvl)
       stream << "abc" << "def"
       res = stream.finish
       expect(Zstd.decompress(res)).to eq('abcdef')
@@ -56,14 +56,26 @@
   if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0')
     describe 'Ractor' do
       it 'should be supported' do
-        r = Ractor.new {
-          stream = Zstd::StreamingCompress.new(5)
+        r = Ractor.new(no_gvl) do |no_gvl|
+          stream = Zstd::StreamingCompress.new(5, no_gvl: no_gvl)
           stream << "abc" << "def"
           res = stream.finish
-        }
+        end
         expect(Zstd.decompress(r.take)).to eq('abcdef')
       end
     end
   end
 end
 
+RSpec.describe Zstd::StreamingCompress do
+  describe "with the global lock" do
+    let(:no_gvl) { false }
+    it_behaves_like "a streaming compressor"
+  end
+
+  describe "without the global lock" do
+    let(:no_gvl) { true }
+    it_behaves_like "a streaming compressor"
+  end
+end
+
diff --git a/spec/zstd-ruby-streaming-decompress_spec.rb b/spec/zstd-ruby-streaming-decompress_spec.rb
@@ -2,7 +2,7 @@
 require 'zstd-ruby'
 require 'securerandom'
 
-RSpec.describe Zstd::StreamingDecompress do
+shared_examples "a streaming decompressor" do
   describe 'streaming decompress' do
     it 'shoud work' do
       # str = SecureRandom.hex(150)
@@ -22,7 +22,7 @@
       # str = SecureRandom.hex(150)
       str = "foo bar buzz" * 100
       cstr = Zstd.compress(str)
-      stream = Zstd::StreamingDecompress.new
+      stream = Zstd::StreamingDecompress.new(no_gvl: no_gvl)
       result = ''
       result << stream.decompress(cstr[0, 5])
       result << stream.decompress(cstr[5, 5])
@@ -35,18 +35,30 @@
   if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0')
     describe 'Ractor' do
       it 'should be supported' do
-        r = Ractor.new {
+        r = Ractor.new(no_gvl) do |no_gvl|
           cstr = Zstd.compress('foo bar buzz')
-          stream = Zstd::StreamingDecompress.new
+          stream = Zstd::StreamingDecompress.new(no_gvl: no_gvl)
           result = ''
           result << stream.decompress(cstr[0, 5])
           result << stream.decompress(cstr[5, 5])
           result << stream.decompress(cstr[10..-1])
           result
-        }
+        end
         expect(r.take).to eq('foo bar buzz')
       end
     end
   end
 end
 
+RSpec.describe Zstd::StreamingDecompress do
+  describe "with the gvl" do
+    let(:no_gvl) { false }
+    it_behaves_like "a streaming decompressor"
+  end
+
+  describe "without the gvl" do
+    let(:no_gvl) { true }
+    it_behaves_like "a streaming decompressor"
+  end
+end
+
