[PATCH] socket: nonblocking sendmsg/recvmsg can be exception-free

[PATCH] socket: nonblocking sendmsg/recvmsg can be exception-free

[Eric Wong]

---
 ext/socket/ancdata.c         | 19 ++++++++++++++++---
 ext/socket/init.c            |  9 ++++++++-
 test/socket/test_nonblock.rb |  4 ++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/ext/socket/ancdata.c b/ext/socket/ancdata.c
index 0a94149..1a8f854 100644
--- a/ext/socket/ancdata.c
+++ b/ext/socket/ancdata.c
@@ -1356,6 +1356,7 @@ rsock_bsock_sendmsg_nonblock(int argc, VALUE *argv, VALUE sock)
 #endif
 
 #if defined(HAVE_RECVMSG)
+static VALUE sym_exception, sym_wait_readable;
 struct recvmsg_args_struct {
     int fd;
     int flags;
@@ -1495,6 +1496,7 @@ bsock_recvmsg_internal(int argc, VALUE *argv, VALUE sock, int nonblock)
     VALUE ret;
     ssize_t ss;
     int request_scm_rights;
+    int ex = 1;
 #if defined(HAVE_STRUCT_MSGHDR_MSG_CONTROL)
     struct cmsghdr *cmh;
     size_t maxctllen;
@@ -1528,8 +1530,12 @@ bsock_recvmsg_internal(int argc, VALUE *argv, VALUE sock, int nonblock)
     grow_buffer = NIL_P(vmaxdatlen) || NIL_P(vmaxctllen);
 
     request_scm_rights = 0;
-    if (!NIL_P(vopts) && RTEST(rb_hash_aref(vopts, ID2SYM(rb_intern("scm_rights")))))
-        request_scm_rights = 1;
+    if (!NIL_P(vopts)) {
+	if (RTEST(rb_hash_aref(vopts, ID2SYM(rb_intern("scm_rights")))))
+	    request_scm_rights = 1;
+	if (nonblock && Qfalse == rb_hash_lookup2(vopts, sym_exception, Qundef))
+	    ex = 0;
+    }
 #if !defined(HAVE_STRUCT_MSGHDR_MSG_CONTROL)
     if (request_scm_rights)
         rb_raise(rb_eNotImpError, "control message for recvmsg is unimplemented");
@@ -1605,8 +1611,10 @@ bsock_recvmsg_internal(int argc, VALUE *argv, VALUE sock, int nonblock)
             rb_io_check_closed(fptr);
             goto retry;
         }
-        if (nonblock && (errno == EWOULDBLOCK || errno == EAGAIN))
+        if (nonblock && (errno == EWOULDBLOCK || errno == EAGAIN)) {
+            if (ex) return sym_wait_readable;
             rb_readwrite_sys_fail(RB_IO_WAIT_READABLE, "recvmsg(2) would block");
+	}
 #if defined(HAVE_STRUCT_MSGHDR_MSG_CONTROL)
         if (!gc_done && (errno == EMFILE || errno == EMSGSIZE)) {
           /*
@@ -1836,4 +1844,9 @@ rsock_init_ancdata(void)
     rb_define_method(rb_cAncillaryData, "ipv6_pktinfo_addr", ancillary_ipv6_pktinfo_addr, 0);
     rb_define_method(rb_cAncillaryData, "ipv6_pktinfo_ifindex", ancillary_ipv6_pktinfo_ifindex, 0);
 #endif
+#if defined(HAVE_RECVMSG)
+#undef rb_intern
+    sym_exception = ID2SYM(rb_intern("exception"));
+    sym_wait_readable = ID2SYM(rb_intern("wait_readable"));
+#endif
 }
diff --git a/ext/socket/init.c b/ext/socket/init.c
index 455652d..eadeba9 100644
--- a/ext/socket/init.c
+++ b/ext/socket/init.c
@@ -188,14 +188,19 @@ rsock_s_recvfrom_nonblock(VALUE sock, int argc, VALUE *argv, enum sock_recv_type
     long slen;
     int fd, flags;
     VALUE addr = Qnil;
+    VALUE opts = Qnil;
     socklen_t len0;
+    int ex = 1;
 
-    rb_scan_args(argc, argv, "11", &len, &flg);
+    rb_scan_args(argc, argv, "11:", &len, &flg, &opts);
 
     if (flg == Qnil) flags = 0;
     else             flags = NUM2INT(flg);
     buflen = NUM2INT(len);
 
+    if (!NIL_P(opts) && Qfalse == rb_hash_lookup2(opts, sym_exception, Qundef))
+	ex = 0;
+
 #ifdef MSG_DONTWAIT
     /* MSG_DONTWAIT avoids the race condition between fcntl and recvfrom.
        It is not portable, though. */
@@ -226,6 +231,8 @@ rsock_s_recvfrom_nonblock(VALUE sock, int argc, VALUE *argv, enum sock_recv_type
 #if defined(EWOULDBLOCK) && EWOULDBLOCK != EAGAIN
 	  case EWOULDBLOCK:
 #endif
+            if (!ex)
+		return sym_wait_readable;
             rb_readwrite_sys_fail(RB_IO_WAIT_READABLE, "recvfrom(2) would block");
 	}
 	rb_sys_fail("recvfrom(2)");
diff --git a/test/socket/test_nonblock.rb b/test/socket/test_nonblock.rb
index 0109d19..00da116 100644
--- a/test/socket/test_nonblock.rb
+++ b/test/socket/test_nonblock.rb
@@ -89,6 +89,8 @@ class TestSocketNonblock < Test::Unit::TestCase
     u2 = UDPSocket.new
     u1.bind("127.0.0.1", 0)
     assert_raise(IO::WaitReadable) { u1.recvfrom_nonblock(100) }
+    assert_raise(IO::WaitReadable) { u1.recvfrom_nonblock(99, exception: true) }
+    assert_equal(:wait_readable, u1.recvfrom_nonblock(100, exception: false))
     assert_raise(IO::WaitReadable, Errno::EINVAL) { u2.recvfrom_nonblock(100) }
     u2.send("aaa", 0, u1.getsockname)
     IO.select [u1]
@@ -115,6 +117,8 @@ class TestSocketNonblock < Test::Unit::TestCase
     u2 = UDPSocket.new
     u1.bind("127.0.0.1", 0)
     assert_raise(IO::WaitReadable) { u1.recv_nonblock(100) }
+    assert_raise(IO::WaitReadable) { u1.recv_nonblock(100, exception: true) }
+    assert_equal(:wait_readable, u1.recv_nonblock(100, exception: false))
     assert_raise(IO::WaitReadable, Errno::EINVAL) { u2.recv_nonblock(100) }
     u2.send("aaa", 0, u1.getsockname)
     IO.select [u1]
-- 
EW