Lock IO::Buffer allocations during copy. (#18489)

Merged-By: ioquatix <samuel@codeotaku.com>

ext/-test-/io_buffer/io_buffer.c

@@ -89,6 +89,74 @@ io_buffer_for_writing_modify_string(VALUE self, VALUE string)
     return rb_io_buffer_for_writing(string, io_buffer_modify_string, string);
 }
 
+static VALUE
+io_buffer_locked_for_reading_callback(const void *base, size_t size, VALUE buffer)
+{
+    VALUE result = rb_ary_new_capa(3);
+
+    rb_ary_push(result, rb_funcall(buffer, rb_intern("locked?"), 0));
+    rb_ary_push(result, SIZET2NUM(size));
+    rb_ary_push(result, size > 0 ? INT2FIX(*(const unsigned char *)base) : Qnil);
+
+    return result;
+}
+
+static VALUE
+io_buffer_locked_for_reading(VALUE self, VALUE buffer)
+{
+    return rb_io_buffer_locked_for_reading(buffer, io_buffer_locked_for_reading_callback, buffer);
+}
+
+static VALUE
+io_buffer_locked_for_reading_raise_callback(const void *base, size_t size, VALUE argument)
+{
+    (void)base;
+    (void)size;
+    (void)argument;
+
+    rb_raise(rb_eRuntimeError, "interrupted");
+}
+
+static VALUE
+io_buffer_locked_for_reading_raise(VALUE self, VALUE buffer)
+{
+    return rb_io_buffer_locked_for_reading(buffer, io_buffer_locked_for_reading_raise_callback, Qnil);
+}
+
+static VALUE
+io_buffer_locked_for_writing_callback(void *base, size_t size, VALUE buffer)
+{
+    VALUE locked = rb_funcall(buffer, rb_intern("locked?"), 0);
+
+    if (size > 0) {
+        *(unsigned char *)base = 'x';
+    }
+
+    return locked;
+}
+
+static VALUE
+io_buffer_locked_for_writing(VALUE self, VALUE buffer)
+{
+    return rb_io_buffer_locked_for_writing(buffer, io_buffer_locked_for_writing_callback, buffer);
+}
+
+static VALUE
+io_buffer_locked_for_writing_raise_callback(void *base, size_t size, VALUE argument)
+{
+    (void)base;
+    (void)size;
+    (void)argument;
+
+    rb_raise(rb_eRuntimeError, "interrupted");
+}
+
+static VALUE
+io_buffer_locked_for_writing_raise(VALUE self, VALUE buffer)
+{
+    return rb_io_buffer_locked_for_writing(buffer, io_buffer_locked_for_writing_raise_callback, Qnil);
+}
+
 static VALUE
 io_buffer_lock(VALUE self, VALUE buffer)
 {
@@ -126,6 +194,10 @@ Init_io_buffer(void)
     rb_define_singleton_method(mIOBuffer, "for_writing_set_string", io_buffer_for_writing_set_string, 2);
     rb_define_singleton_method(mIOBuffer, "for_writing_readonly?", io_buffer_for_writing_readonly_p, 1);
     rb_define_singleton_method(mIOBuffer, "for_writing_modify_string", io_buffer_for_writing_modify_string, 1);
+    rb_define_singleton_method(mIOBuffer, "locked_for_reading", io_buffer_locked_for_reading, 1);
+    rb_define_singleton_method(mIOBuffer, "locked_for_reading_raise", io_buffer_locked_for_reading_raise, 1);
+    rb_define_singleton_method(mIOBuffer, "locked_for_writing", io_buffer_locked_for_writing, 1);
+    rb_define_singleton_method(mIOBuffer, "locked_for_writing_raise", io_buffer_locked_for_writing_raise, 1);
     rb_define_singleton_method(mIOBuffer, "lock", io_buffer_lock, 1);
     rb_define_singleton_method(mIOBuffer, "unlock", io_buffer_unlock, 1);
     rb_define_singleton_method(mIOBuffer, "new_locked", io_buffer_new_locked, 1);

include/ruby/io/buffer.h

@@ -104,6 +104,18 @@ enum rb_io_buffer_flags rb_io_buffer_get_bytes(VALUE self, void **base, size_t *
 void rb_io_buffer_get_bytes_for_reading(VALUE self, const void **base, size_t *size);
 void rb_io_buffer_get_bytes_for_writing(VALUE self, void **base, size_t *size);
 
+// Lock the backing allocation, invoke the callback with its readable bytes,
+// and automatically unlock it when the callback returns or raises. The bytes
+// are only valid for the duration of the callback. This protects the lifetime
+// of the allocation; it does not provide synchronization for its contents.
+VALUE rb_io_buffer_locked_for_reading(VALUE self, VALUE (*callback)(const void *base, size_t size, VALUE argument), VALUE argument);
+
+// Lock the backing allocation, invoke the callback with its writable bytes,
+// and automatically unlock it when the callback returns or raises. The bytes
+// are only valid for the duration of the callback. This protects the lifetime
+// of the allocation; it does not provide synchronization for its contents.
+VALUE rb_io_buffer_locked_for_writing(VALUE self, VALUE (*callback)(void *base, size_t size, VALUE argument), VALUE argument);
+
 VALUE rb_io_buffer_transfer(VALUE self);
 void rb_io_buffer_resize(VALUE self, size_t size);
 void rb_io_buffer_clear(VALUE self, uint8_t value, size_t offset, size_t length);

io_buffer.c

@@ -1673,6 +1673,73 @@ rb_io_buffer_locked_ensure(VALUE self)
     return Qnil;
 }
 
+struct io_buffer_readable_bytes_arguments {
+    VALUE self;
+    VALUE (*callback)(const void *base, size_t size, VALUE argument);
+    VALUE argument;
+};
+
+static VALUE
+io_buffer_readable_bytes_call(VALUE _arguments)
+{
+    struct io_buffer_readable_bytes_arguments *arguments = (void *)_arguments;
+
+    const void *base;
+    size_t size;
+    rb_io_buffer_get_bytes_for_reading(arguments->self, &base, &size);
+
+    return arguments->callback(base, size, arguments->argument);
+}
+
+VALUE
+rb_io_buffer_locked_for_reading(VALUE self, VALUE (*callback)(const void *base, size_t size, VALUE argument), VALUE argument)
+{
+    struct rb_io_buffer *buffer = get_io_buffer(self);
+    io_buffer_validate_for_reading(buffer);
+
+    struct io_buffer_readable_bytes_arguments arguments = {
+        .self = self,
+        .callback = callback,
+        .argument = argument,
+    };
+
+    rb_io_buffer_lock(self);
+    return rb_ensure(io_buffer_readable_bytes_call, (VALUE)&arguments, rb_io_buffer_locked_ensure, self);
+}
+
+struct io_buffer_writable_bytes_arguments {
+    VALUE self;
+    VALUE (*callback)(void *base, size_t size, VALUE argument);
+    VALUE argument;
+};
+
+static VALUE
+io_buffer_writable_bytes_call(VALUE _arguments)
+{
+    struct io_buffer_writable_bytes_arguments *arguments = (void *)_arguments;
+
+    void *base;
+    size_t size;
+    rb_io_buffer_get_bytes_for_writing(arguments->self, &base, &size);
+
+    return arguments->callback(base, size, arguments->argument);
+}
+
+VALUE
+rb_io_buffer_locked_for_writing(VALUE self, VALUE (*callback)(void *base, size_t size, VALUE argument), VALUE argument)
+{
+    get_io_buffer_for_writing(self);
+
+    struct io_buffer_writable_bytes_arguments arguments = {
+        .self = self,
+        .callback = callback,
+        .argument = argument,
+    };
+
+    rb_io_buffer_lock(self);
+    return rb_ensure(io_buffer_writable_bytes_call, (VALUE)&arguments, rb_io_buffer_locked_ensure, self);
+}
+
 /*
  *  call-seq: locked { ... }
  *
@@ -2891,13 +2958,11 @@ io_buffer_memmove_unblock(void *data)
 }
 
 static void
-io_buffer_memmove(struct rb_io_buffer *buffer, size_t offset, const void *source_base, size_t source_offset, size_t source_size, size_t length)
+io_buffer_memmove(void *base, size_t size, size_t offset, const void *source_base, size_t source_offset, size_t source_size, size_t length)
 {
-    void *base;
-    size_t size;
-    io_buffer_get_bytes_for_writing(buffer, &base, &size);
-
-    io_buffer_validate_range(buffer, offset, length);
+    if (size_sum_is_bigger_than(offset, length, size)) {
+        rb_raise(rb_eArgError, "Specified offset+length is bigger than the buffer size!");
+    }
 
     if (size_sum_is_bigger_than(source_offset, length, source_size)) {
         rb_raise(rb_eArgError, "The computed source range exceeds the size of the source buffer!");
@@ -2920,45 +2985,114 @@ io_buffer_memmove(struct rb_io_buffer *buffer, size_t offset, const void *source
     }
 }
 
-// (offset, length, source_offset) -> length
-static VALUE
-io_buffer_copy_from(struct rb_io_buffer *buffer, const void *source_base, size_t source_size, int argc, VALUE *argv)
+static void
+io_buffer_extract_copy_arguments(size_t source_size, int argc, VALUE *argv, size_t *offset, size_t *length, size_t *source_offset)
 {
-    size_t offset = 0;
-    size_t length;
-    size_t source_offset;
-
     // The offset we copy into the buffer:
     if (argc >= 1) {
-        offset = io_buffer_extract_offset(argv[0]);
+        *offset = io_buffer_extract_offset(argv[0]);
+    }
+    else {
+        *offset = 0;
     }
 
     // The offset we start from within the string:
     if (argc >= 3) {
-        source_offset = io_buffer_extract_offset(argv[2]);
+        *source_offset = io_buffer_extract_offset(argv[2]);
 
-        if (source_offset > source_size) {
+        if (*source_offset > source_size) {
             rb_raise(rb_eArgError, "The given source offset is bigger than the source itself!");
         }
     }
     else {
-        source_offset = 0;
+        *source_offset = 0;
     }
 
     // The length we are going to copy:
     if (argc >= 2 && !RB_NIL_P(argv[1])) {
-        length = io_buffer_extract_length(argv[1]);
+        *length = io_buffer_extract_length(argv[1]);
     }
     else {
         // Default to the source offset -> source size:
-        length = source_size - source_offset;
+        *length = source_size - *source_offset;
     }
+}
 
-    io_buffer_memmove(buffer, offset, source_base, source_offset, source_size, length);
+// (offset, length, source_offset) -> length
+static VALUE
+io_buffer_copy_from(struct rb_io_buffer *buffer, const void *source_base, size_t source_size, int argc, VALUE *argv)
+{
+    size_t offset, length, source_offset;
+    io_buffer_extract_copy_arguments(source_size, argc, argv, &offset, &length, &source_offset);
+
+    void *base;
+    size_t size;
+    io_buffer_get_bytes_for_writing(buffer, &base, &size);
+
+    io_buffer_memmove(base, size, offset, source_base, source_offset, source_size, length);
 
     return SIZET2NUM(length);
 }
 
+struct io_buffer_copy_arguments {
+    VALUE destination;
+    const void *source_base;
+    size_t source_size;
+    int argc;
+    VALUE *argv;
+};
+
+// This is the innermost callback for IO::Buffer#copy. At this point the source
+// is locked for reading and the destination is locked for writing, so both
+// pointers and sizes remain valid while arguments are extracted, ranges are
+// validated, and memmove potentially releases the GVL.
+static VALUE
+io_buffer_copy_to(void *base, size_t size, VALUE _arguments)
+{
+    struct io_buffer_copy_arguments *arguments = (void *)_arguments;
+
+    size_t offset, length, source_offset;
+    io_buffer_extract_copy_arguments(arguments->source_size, arguments->argc, arguments->argv, &offset, &length, &source_offset);
+
+    io_buffer_memmove(base, size, offset, arguments->source_base, source_offset, arguments->source_size, length);
+
+    return SIZET2NUM(length);
+}
+
+// This callback runs while the source is locked for reading. Retain its bytes
+// in the callback arguments, then enter the destination's writable scope. The
+// source scope remains active until that nested scope returns.
+static VALUE
+io_buffer_copy_from_readable(const void *base, size_t size, VALUE _arguments)
+{
+    struct io_buffer_copy_arguments *arguments = (void *)_arguments;
+
+    arguments->source_base = base;
+    arguments->source_size = size;
+
+    return rb_io_buffer_locked_for_writing(arguments->destination, io_buffer_copy_to, _arguments);
+}
+
+static VALUE
+io_buffer_initialize_copy_from(const void *base, size_t size, VALUE self)
+{
+    struct rb_io_buffer *buffer = get_io_buffer(self);
+
+    io_buffer_initialize(self, buffer, NULL, size, io_flags_for_size(size), Qnil);
+
+    struct io_buffer_copy_arguments arguments = {
+        .destination = self,
+        .source_base = base,
+        .source_size = size,
+        .argc = 0,
+        .argv = NULL,
+    };
+
+    // The source remains locked by the outer readable scope while the newly
+    // initialized destination is locked and populated by io_buffer_copy_to.
+    return rb_io_buffer_locked_for_writing(self, io_buffer_copy_to, (VALUE)&arguments);
+}
+
 /*
  *  call-seq:
  *    dup -> io_buffer
@@ -2979,18 +3113,7 @@ io_buffer_copy_from(struct rb_io_buffer *buffer, const void *source_base, size_t
 static VALUE
 rb_io_buffer_initialize_copy(VALUE self, VALUE source)
 {
-    struct rb_io_buffer *buffer = get_io_buffer(self);
-
-    const void *source_base;
-    size_t source_size;
-
-    rb_io_buffer_get_bytes_for_reading(source, &source_base, &source_size);
-
-    io_buffer_initialize(self, buffer, NULL, source_size, io_flags_for_size(source_size), Qnil);
-
-    VALUE result = io_buffer_copy_from(buffer, source_base, source_size, 0, NULL);
-    RB_GC_GUARD(source);
-    return result;
+    return rb_io_buffer_locked_for_reading(source, io_buffer_initialize_copy_from, self);
 }
 
 /*
@@ -3066,17 +3189,19 @@ io_buffer_copy(int argc, VALUE *argv, VALUE self)
 {
     rb_check_arity(argc, 1, 4);
 
-    struct rb_io_buffer *buffer = get_io_buffer(self);
-
     VALUE source = argv[0];
-    const void *source_base;
-    size_t source_size;
+    struct io_buffer_copy_arguments arguments = {
+        .destination = self,
+        .argc = argc-1,
+        .argv = argv+1,
+    };
 
-    rb_io_buffer_get_bytes_for_reading(source, &source_base, &source_size);
-
-    VALUE result = io_buffer_copy_from(buffer, source_base, source_size, argc-1, argv+1);
-    RB_GC_GUARD(source);
-    return result;
+    // Lock the source first, then io_buffer_copy_from_readable nests the
+    // destination lock. The scoped helpers use rb_ensure, so the destination
+    // is unlocked before the source on both normal and exceptional returns.
+    // If both buffers share an allocation, its reference-counted lock is
+    // acquired and released twice.
+    return rb_io_buffer_locked_for_reading(source, io_buffer_copy_from_readable, (VALUE)&arguments);
 }
 
 /*

test/ruby/test_io_buffer.rb

@@ -94,6 +94,52 @@ class TestIOBuffer < Test::Unit::TestCase
     assert_equal "hello!", string
   end
 
+  def test_internal_locked_for_reading
+    buffer = IO::Buffer.for("hello")
+
+    assert_equal [true, 5, "h".ord], Bug::IOBuffer.locked_for_reading(buffer)
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_locked_for_reading_unlocks_after_callback_exception
+    buffer = IO::Buffer.for("hello")
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.locked_for_reading_raise(buffer)
+    end
+
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_locked_for_writing
+    buffer = IO::Buffer.new(5)
+    buffer.set_string("hello")
+
+    assert_equal true, Bug::IOBuffer.locked_for_writing(buffer)
+    assert_equal "xello", buffer.get_string
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_locked_for_writing_rejects_readonly_buffer
+    buffer = IO::Buffer.for("hello")
+
+    assert_raise(IO::Buffer::AccessError) do
+      Bug::IOBuffer.locked_for_writing(buffer)
+    end
+
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_locked_for_writing_unlocks_after_callback_exception
+    buffer = IO::Buffer.new(5)
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.locked_for_writing_raise(buffer)
+    end
+
+    refute_predicate buffer, :locked?
+  end
+
   def test_endian
     assert_equal 4, IO::Buffer::LITTLE_ENDIAN
     assert_equal 8, IO::Buffer::BIG_ENDIAN