Merge pull request #831 from byroot/reduce-buffering-do-write-overhead

Reduce OpenSSL::Buffering#do_write overhead
ruby · Jan 14, 2025 · a3e0c16 · a3e0c16
2 parents 2733bea + 28f2901
commit a3e0c16
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 30 deletions.
diff --git a/ext/openssl/ossl_ssl.c b/ext/openssl/ossl_ssl.c
@@ -2054,28 +2054,32 @@ ossl_ssl_read_nonblock(int argc, VALUE *argv, VALUE self)
 }
 
 static VALUE
-ossl_ssl_write_internal(VALUE self, VALUE str, VALUE opts)
+ossl_ssl_write_internal_safe(VALUE _args)
 {
+    VALUE *args = (VALUE*)_args;
+    VALUE self = args[0];
+    VALUE str = args[1];
+    VALUE opts = args[2];
+
     SSL *ssl;
     rb_io_t *fptr;
     int num, nonblock = opts != Qfalse;
-    VALUE tmp, cb_state;
+    VALUE cb_state;
 
     GetSSL(self, ssl);
     if (!ssl_started(ssl))
         rb_raise(eSSLError, "SSL session is not started yet");
 
-    tmp = rb_str_new_frozen(StringValue(str));
     VALUE io = rb_attr_get(self, id_i_io);
     GetOpenFile(io, fptr);
 
     /* SSL_write(3ssl) manpage states num == 0 is undefined */
-    num = RSTRING_LENINT(tmp);
+    num = RSTRING_LENINT(str);
     if (num == 0)
         return INT2FIX(0);
 
     for (;;) {
-        int nwritten = SSL_write(ssl, RSTRING_PTR(tmp), num);
+        int nwritten = SSL_write(ssl, RSTRING_PTR(str), num);
 
         cb_state = rb_attr_get(self, ID_callback_state);
         if (!NIL_P(cb_state)) {
@@ -2116,6 +2120,29 @@ ossl_ssl_write_internal(VALUE self, VALUE str, VALUE opts)
     }
 }
 
+
+static VALUE
+ossl_ssl_write_internal(VALUE self, VALUE str, VALUE opts)
+{
+    VALUE args[3] = {self, str, opts};
+    int state;
+    str = StringValue(str);
+
+    int frozen = RB_OBJ_FROZEN(str);
+    if (!frozen) {
+        str = rb_str_locktmp(str);
+    }
+    VALUE result = rb_protect(ossl_ssl_write_internal_safe, (VALUE)args, &state);
+    if (!frozen) {
+        rb_str_unlocktmp(str);
+    }
+
+    if (state) {
+        rb_jump_tag(state);
+    }
+    return result;
+}
+
 /*
  * call-seq:
  *    ssl.syswrite(string) => Integer

diff --git a/lib/openssl/buffering.rb b/lib/openssl/buffering.rb
@@ -24,25 +24,21 @@ module OpenSSL::Buffering
 
   # A buffer which will retain binary encoding.
   class Buffer < String
-    BINARY = Encoding::BINARY
-
-    def initialize
-      super
-
-      force_encoding(BINARY)
-    end
+    unless String.method_defined?(:append_as_bytes)
+      alias_method :_append, :<<
+      def append_as_bytes(string)
+        if string.encoding == Encoding::BINARY
+          _append(string)
+        else
+          _append(string.b)
+        end
 
-    def << string
-      if string.encoding == BINARY
-        super(string)
-      else
-        super(string.b)
+        self
       end
-
-      return self
     end
 
-    alias concat <<
+    undef_method :concat
+    undef_method :<<
   end
 
   ##
@@ -77,7 +73,7 @@ def initialize(*)
 
   def fill_rbuff
     begin
-      @rbuffer << self.sysread(BLOCK_SIZE)
+      @rbuffer.append_as_bytes(self.sysread(BLOCK_SIZE))
     rescue Errno::EAGAIN
       retry
     rescue EOFError
@@ -352,22 +348,32 @@ def eof?
 
   def do_write(s)
     @wbuffer = Buffer.new unless defined? @wbuffer
-    @wbuffer << s
-    @wbuffer.force_encoding(Encoding::BINARY)
+    @wbuffer.append_as_bytes(s)
+
     @sync ||= false
-    buffer_size = @wbuffer.size
+    buffer_size = @wbuffer.bytesize
     if @sync or buffer_size > BLOCK_SIZE
       nwrote = 0
       begin
         while nwrote < buffer_size do
           begin
-            nwrote += syswrite(@wbuffer[nwrote, buffer_size - nwrote])
+            chunk = if nwrote > 0
+              @wbuffer.byteslice(nwrote, @wbuffer.bytesize)
+            else
+              @wbuffer
+            end
+
+            nwrote += syswrite(chunk)
           rescue Errno::EAGAIN
             retry
           end
         end
       ensure
-        @wbuffer[0, nwrote] = ""
+        if nwrote < @wbuffer.bytesize
+          @wbuffer[0, nwrote] = ""
+        else
+          @wbuffer.clear
+        end
       end
     end
   end
@@ -444,10 +450,10 @@ def <<(s)
   def puts(*args)
     s = Buffer.new
     if args.empty?
-      s << "\n"
+      s.append_as_bytes("\n")
     end
     args.each{|arg|
-      s << arg.to_s
+      s.append_as_bytes(arg.to_s)
       s.sub!(/(?<!\n)\z/, "\n")
     }
     do_write(s)
@@ -461,7 +467,7 @@ def puts(*args)
 
   def print(*args)
     s = Buffer.new
-    args.each{ |arg| s << arg.to_s }
+    args.each{ |arg| s.append_as_bytes(arg.to_s) }
     do_write(s)
     nil
   end

diff --git a/test/openssl/test_buffering.rb b/test/openssl/test_buffering.rb
@@ -31,7 +31,7 @@ def sysread(size)
     end
 
     def syswrite(str)
-      @io << str
+      @io.append_as_bytes(str)
       str.size
     end
   end