socket: allow explicit buffer for recv and recv_nonblock

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

* ext/socket/basicsocket.c (bsock_recv): document outbuf
* ext/socket/unixsocket.c (unix_recvfrom): ditto
* ext/socket/init.c (rsock_strbuf, recvfrom_locktmp): new functions
  (rsock_s_recvfrom): support destination buffer as 3rd arg
  (rsock_s_recvfrom_nonblock): ditto
* string.c (rb_str_locktmp_ensure): export for internal ext
* test/socket/test_nonblock.rb: test recv_nonblock
* test/socket/test_unix.rb: test recv
  [ruby-core:69543] [Feature #11242]

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

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@50912 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

3 files changed, 44 insertions, 19 deletions

diff --git a/ext/socket/basicsocket.c b/ext/socket/basicsocket.c
index 4c793263ef..fbe5167c15 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 5f0d445c70..ab6b52aab3 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 9a4c2cfc01..9003529042 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)