[PATCH] socket: allow explicit buffer for recv and recv_nonblock

[PATCH] socket: allow explicit buffer for recv and recv_nonblock

[Eric Wong]

This reduces GC overhead and makes the API more consistent
with IO#read and IO#read_nonblock.

Benchmark results:

             user     system      total        real
alloc    0.130000   0.280000   0.410000 (  0.420656)
extbuf   0.100000   0.220000   0.320000 (  0.318708)

-------------------8<--------------------
require 'socket'
require 'benchmark'
nr = 100000
msg = ' ' * 16384
size = msg.bytesize
buf = ' ' * size
UNIXSocket.pair(:DGRAM) do |a, b|
  Benchmark.bmbm do |x|
    x.report('alloc') do
      nr.times do
        b.send(msg, 0)
        a.recv(size, 0)
      end
    end

    x.report('extbuf') do
      nr.times do
        b.send(msg, 0)
        a.recv(size, 0, buf)
      end
    end
  end
end
---
 ext/socket/basicsocket.c     |  6 +++--
 ext/socket/init.c            | 52 ++++++++++++++++++++++++++++++--------------
 ext/socket/unixsocket.c      |  5 ++++-
 string.c                     |  2 +-
 test/socket/test_nonblock.rb | 10 +++++++++
 test/socket/test_unix.rb     |  7 ++++++
 6 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/ext/socket/basicsocket.c b/ext/socket/basicsocket.c
index 4c79326..fbe5167 100644
--- a/ext/socket/basicsocket.c
+++ b/ext/socket/basicsocket.c
@@ -615,8 +615,7 @@ bsock_do_not_reverse_lookup_set(VALUE sock, VALUE state)
 
 /*
  * call-seq:
- *   basicsocket.recv(maxlen) => mesg
- *   basicsocket.recv(maxlen, flags) => mesg
+ *   basicsocket.recv(maxlen[, flags[, outbuf]]) => mesg
  *
  * Receives a message.
  *
@@ -624,6 +623,9 @@ bsock_do_not_reverse_lookup_set(VALUE sock, VALUE state)
  *
  * _flags_ should be a bitwise OR of Socket::MSG_* constants.
  *
+ * _outbuf_ will contain only the received data after the method call
+ * even if it is not empty at the beginning.
+ *
  *   UNIXSocket.pair {|s1, s2|
  *     s1.puts "Hello World"
  *     p s2.recv(4)                     #=> "Hell"
diff --git a/ext/socket/init.c b/ext/socket/init.c
index 5f0d445..ddfff57 100644
--- a/ext/socket/init.c
+++ b/ext/socket/init.c
@@ -108,21 +108,48 @@ recvfrom_blocking(void *data)
     return (VALUE)ret;
 }
 
+static VALUE
+rsock_strbuf(VALUE str, long buflen)
+{
+    long len;
+
+    if (NIL_P(str)) return rb_tainted_str_new(0, buflen);
+
+    StringValue(str);
+    len = RSTRING_LEN(str);
+    if (len >= buflen) {
+	rb_str_modify(str);
+    } else {
+	rb_str_modify_expand(str, buflen - len);
+    }
+    rb_str_set_len(str, buflen);
+    return str;
+}
+
+static VALUE
+recvfrom_locktmp(VALUE v)
+{
+  struct recvfrom_arg *arg = (struct recvfrom_arg *)v;
+
+  return rb_thread_io_blocking_region(recvfrom_blocking, arg, arg->fd);
+}
+
 VALUE
 rsock_s_recvfrom(VALUE sock, int argc, VALUE *argv, enum sock_recv_type from)
 {
     rb_io_t *fptr;
-    VALUE str, klass;
+    VALUE str;
     struct recvfrom_arg arg;
     VALUE len, flg;
     long buflen;
     long slen;
 
-    rb_scan_args(argc, argv, "11", &len, &flg);
+    rb_scan_args(argc, argv, "12", &len, &flg, &str);
 
     if (flg == Qnil) arg.flags = 0;
     else             arg.flags = NUM2INT(flg);
     buflen = NUM2INT(len);
+    str = rsock_strbuf(str, buflen);
 
     GetOpenFile(sock, fptr);
     if (rb_io_read_pending(fptr)) {
@@ -130,24 +157,18 @@ rsock_s_recvfrom(VALUE sock, int argc, VALUE *argv, enum sock_recv_type from)
     }
     arg.fd = fptr->fd;
     arg.alen = (socklen_t)sizeof(arg.buf);
-
-    arg.str = str = rb_tainted_str_new(0, buflen);
-    klass = RBASIC(str)->klass;
-    rb_obj_hide(str);
+    arg.str = str;
 
     while (rb_io_check_closed(fptr),
 	   rsock_maybe_wait_fd(arg.fd),
-	   (slen = BLOCKING_REGION_FD(recvfrom_blocking, &arg)) < 0) {
+	   (slen = (long)rb_str_locktmp_ensure(str, recvfrom_locktmp,
+	                                       (VALUE)&arg)) < 0) {
         if (!rb_io_wait_readable(fptr->fd)) {
             rb_sys_fail("recvfrom(2)");
         }
-	if (RBASIC(str)->klass || RSTRING_LEN(str) != buflen) {
-	    rb_raise(rb_eRuntimeError, "buffer string modified");
-	}
     }
 
-    rb_obj_reveal(str, klass);
-    if (slen < RSTRING_LEN(str)) {
+    if (slen != RSTRING_LEN(str)) {
 	rb_str_set_len(str, slen);
     }
     rb_obj_taint(str);
@@ -191,11 +212,12 @@ rsock_s_recvfrom_nonblock(VALUE sock, int argc, VALUE *argv, enum sock_recv_type
     VALUE opts = Qnil;
     socklen_t len0;
 
-    rb_scan_args(argc, argv, "11:", &len, &flg, &opts);
+    rb_scan_args(argc, argv, "12:", &len, &flg, &str, &opts);
 
     if (flg == Qnil) flags = 0;
     else             flags = NUM2INT(flg);
     buflen = NUM2INT(len);
+    str = rsock_strbuf(str, buflen);
 
 #ifdef MSG_DONTWAIT
     /* MSG_DONTWAIT avoids the race condition between fcntl and recvfrom.
@@ -209,8 +231,6 @@ rsock_s_recvfrom_nonblock(VALUE sock, int argc, VALUE *argv, enum sock_recv_type
     }
     fd = fptr->fd;
 
-    str = rb_tainted_str_new(0, buflen);
-
     rb_io_check_closed(fptr);
 
     if (!MSG_DONTWAIT_RELIABLE)
@@ -233,7 +253,7 @@ rsock_s_recvfrom_nonblock(VALUE sock, int argc, VALUE *argv, enum sock_recv_type
 	}
 	rb_sys_fail("recvfrom(2)");
     }
-    if (slen < RSTRING_LEN(str)) {
+    if (slen != RSTRING_LEN(str)) {
 	rb_str_set_len(str, slen);
     }
     rb_obj_taint(str);
diff --git a/ext/socket/unixsocket.c b/ext/socket/unixsocket.c
index 9a4c2cf..9003529 100644
--- a/ext/socket/unixsocket.c
+++ b/ext/socket/unixsocket.c
@@ -131,7 +131,7 @@ unix_path(VALUE sock)
 
 /*
  * call-seq:
- *   unixsocket.recvfrom(maxlen [, flags]) => [mesg, unixaddress]
+ *   unixsocket.recvfrom(maxlen [, flags[, outbuf]) => [mesg, unixaddress]
  *
  * Receives a message via _unixsocket_.
  *
@@ -139,6 +139,9 @@ unix_path(VALUE sock)
  *
  * _flags_ should be a bitwise OR of Socket::MSG_* constants.
  *
+ * _outbuf_ will contain only the received data after the method call
+ * even if it is not empty at the beginning.
+ *
  *   s1 = Socket.new(:UNIX, :DGRAM, 0)
  *   s1_ai = Addrinfo.unix("/tmp/sock1")
  *   s1.bind(s1_ai)
diff --git a/string.c b/string.c
index ca08540..4d2944c 100644
--- a/string.c
+++ b/string.c
@@ -2107,7 +2107,7 @@ rb_str_unlocktmp(VALUE str)
     return str;
 }
 
-VALUE
+RUBY_FUNC_EXPORTED VALUE
 rb_str_locktmp_ensure(VALUE str, VALUE (*func)(VALUE), VALUE arg)
 {
     rb_str_locktmp(str);
diff --git a/test/socket/test_nonblock.rb b/test/socket/test_nonblock.rb
index 4e14c93..93f463d 100644
--- a/test/socket/test_nonblock.rb
+++ b/test/socket/test_nonblock.rb
@@ -128,6 +128,16 @@ class TestSocketNonblock < Test::Unit::TestCase
     }
     mesg = u1.recv_nonblock(100)
     assert_equal("", mesg)
+
+    buf = "short"
+    out = "hello world" * 4
+    out.freeze
+    u2.send(out, 0, u1.getsockname)
+    IO.select [u1]
+    rv = u1.recv_nonblock(100, 0, buf)
+    assert_equal rv.object_id, buf.object_id
+    assert_equal out, rv
+    assert_equal out, buf
   ensure
     u1.close if u1
     u2.close if u2
diff --git a/test/socket/test_unix.rb b/test/socket/test_unix.rb
index a10f2e5..344dcc9 100644
--- a/test/socket/test_unix.rb
+++ b/test/socket/test_unix.rb
@@ -385,6 +385,13 @@ class TestSocket_UNIXSocket < Test::Unit::TestCase
     assert_equal("", s1.recv(10))
     assert_equal("", s1.recv(10))
     assert_raise(IO::EAGAINWaitReadable) { s1.recv_nonblock(10) }
+
+    buf = ""
+    s2.send("BBBBBB", 0)
+    sleep 0.1
+    rv = s1.recv(100, 0, buf)
+    assert_equal buf.object_id, rv.object_id
+    assert_equal "BBBBBB", rv
   ensure
     s1.close if s1
     s2.close if s2
-- 
EW