[PATCH] connect_nonblock supports "exception: false"

[PATCH] connect_nonblock supports "exception: false"

[Eric Wong]

This is for consistency with accept_nonblock arguments and gives a
minor speedup from avoiding exceptions:

Results:

default            0.050000   0.100000   0.150000 (  0.151307)
exception: false   0.030000   0.080000   0.110000 (  0.108840)

----------------------------8<-----------------------
require 'socket'
require 'benchmark'
require 'io/wait'
require 'tmpdir'

host = "127.#{rand 255}.#{rand 255}.#{rand(254) + 1}"

serv = TCPServer.new(host, 0)

nr = 5000

addr = serv.getsockname
pid = fork do
  begin
    serv.accept.close
  rescue => e
    warn "#$$: #{e.message} (#{e.class})"
  end while true
end
at_exit { Process.kill(:TERM, pid) }
serv.close

Benchmark.bmbm do |x|
  x.report("default") do
    nr.times do
      s = Socket.new(:INET, :STREAM)
      s.setsockopt(:SOL_SOCKET, :SO_REUSEADDR, 1)
      begin
        s.connect_nonblock(addr)
      rescue IO::WaitWritable
        s.wait_writable
      end
      s.close
    end
  end
  x.report("exception: false") do
    nr.times do
      s = Socket.new(:INET, :STREAM)
      s.setsockopt(:SOL_SOCKET, :SO_REUSEADDR, 1)
      case s.connect_nonblock(addr, exception: false)
      when :wait_writable
        s.wait_writable
      end
      s.close
    end
  end
end
---
 ext/openssl/ossl_ssl.c       | 34 ++++++++++++++++++++++------------
 ext/socket/socket.c          | 26 ++++++++++++++++++++++----
 test/openssl/test_pair.rb    | 24 ++++++++++++++++++++++--
 test/socket/test_nonblock.rb | 23 +++++++++++++++++++++++
 4 files changed, 89 insertions(+), 18 deletions(-)

diff --git a/ext/openssl/ossl_ssl.c b/ext/openssl/ossl_ssl.c
index 1a67d35..2e12835 100644
--- a/ext/openssl/ossl_ssl.c
+++ b/ext/openssl/ossl_ssl.c
@@ -1330,9 +1330,17 @@ ossl_ssl_connect(VALUE self)
     return ossl_start_ssl(self, SSL_connect, "SSL_connect", 0, 0);
 }
 
+static int
+get_no_exception(VALUE opts)
+{
+    if (!NIL_P(opts) && Qfalse == rb_hash_lookup2(opts, sym_exception, Qundef))
+	return 1;
+    return 0;
+}
+
 /*
  * call-seq:
- *    ssl.connect_nonblock => self
+ *    ssl.connect_nonblock([options]) => self
  *
  * Initiates the SSL/TLS handshake as a client in non-blocking manner.
  *
@@ -1347,12 +1355,22 @@ ossl_ssl_connect(VALUE self)
  *     retry
  *   end
  *
+ * By specifying `exception: false`, the options hash allows you to indicate
+ * that connect_nonblock should not raise an IO::WaitReadable or
+ * IO::WaitWritable exception, but return the symbol :wait_readable or
+ * :wait_writable instead.
  */
 static VALUE
-ossl_ssl_connect_nonblock(VALUE self)
+ossl_ssl_connect_nonblock(int argc, VALUE *argv, VALUE self)
 {
+    int no_exception;
+    VALUE opts = Qnil;
+
+    rb_scan_args(argc, argv, "0:", &opts);
+    no_exception = get_no_exception(opts);
+
     ossl_ssl_setup(self);
-    return ossl_start_ssl(self, SSL_connect, "SSL_connect", 1, 0);
+    return ossl_start_ssl(self, SSL_connect, "SSL_connect", 1, no_exception);
 }
 
 /*
@@ -1369,14 +1387,6 @@ ossl_ssl_accept(VALUE self)
     return ossl_start_ssl(self, SSL_accept, "SSL_accept", 0, 0);
 }
 
-static int
-get_no_exception(VALUE opts)
-{
-    if (!NIL_P(opts) && Qfalse == rb_hash_lookup2(opts, sym_exception, Qundef))
-	return 1;
-    return 0;
-}
-
 /*
  * call-seq:
  *    ssl.accept_nonblock([options]) => self
@@ -2235,7 +2245,7 @@ Init_ossl_ssl(void)
     rb_define_alias(cSSLSocket, "to_io", "io");
     rb_define_method(cSSLSocket, "initialize", ossl_ssl_initialize, -1);
     rb_define_method(cSSLSocket, "connect",    ossl_ssl_connect, 0);
-    rb_define_method(cSSLSocket, "connect_nonblock",    ossl_ssl_connect_nonblock, 0);
+    rb_define_method(cSSLSocket, "connect_nonblock",    ossl_ssl_connect_nonblock, -1);
     rb_define_method(cSSLSocket, "accept",     ossl_ssl_accept, 0);
     rb_define_method(cSSLSocket, "accept_nonblock", ossl_ssl_accept_nonblock, -1);
     rb_define_method(cSSLSocket, "sysread",    ossl_ssl_read, -1);
diff --git a/ext/socket/socket.c b/ext/socket/socket.c
index bb18409..a514b9a 100644
--- a/ext/socket/socket.c
+++ b/ext/socket/socket.c
@@ -10,6 +10,8 @@
 
 #include "rubysocket.h"
 
+static VALUE sym_exception, sym_wait_writable;
+
 static VALUE sock_s_unpack_sockaddr_in(VALUE, VALUE);
 
 void
@@ -440,7 +442,7 @@ sock_connect(VALUE sock, VALUE addr)
 
 /*
  * call-seq:
- *   socket.connect_nonblock(remote_sockaddr) => 0
+ *   socket.connect_nonblock(remote_sockaddr, [options]) => 0
  *
  * Requests a connection to be made on the given +remote_sockaddr+ after
  * O_NONBLOCK is set for the underlying file descriptor.
@@ -477,24 +479,36 @@ sock_connect(VALUE sock, VALUE addr)
  * it is extended by IO::WaitWritable.
  * So IO::WaitWritable can be used to rescue the exceptions for retrying connect_nonblock.
  *
+ * By specifying `exception: false`, the options hash allows you to indicate
+ * that connect_nonblock should not raise an IO::WaitWritable exception, but
+ * return the symbol :wait_writable instead.
+ *
  * === See
  * * Socket#connect
  */
 static VALUE
-sock_connect_nonblock(VALUE sock, VALUE addr)
+sock_connect_nonblock(int argc, VALUE *argv, VALUE sock)
 {
+    VALUE addr;
+    VALUE opts = Qnil;
     VALUE rai;
     rb_io_t *fptr;
     int n;
 
+    rb_scan_args(argc, argv, "1:", &addr, &opts);
     SockAddrStringValueWithAddrinfo(addr, rai);
     addr = rb_str_new4(addr);
     GetOpenFile(sock, fptr);
     rb_io_set_nonblock(fptr);
     n = connect(fptr->fd, (struct sockaddr*)RSTRING_PTR(addr), RSTRING_SOCKLEN(addr));
     if (n < 0) {
-        if (errno == EINPROGRESS)
+        if (errno == EINPROGRESS) {
+            if (!NIL_P(opts) &&
+                    Qfalse == rb_hash_lookup2(opts, sym_exception, Qundef)) {
+                return sym_wait_writable;
+            }
             rb_readwrite_sys_fail(RB_IO_WAIT_WRITABLE, "connect(2) would block");
+	}
 	rsock_sys_fail_raddrinfo_or_sockaddr("connect(2)", addr, rai);
     }
 
@@ -2158,7 +2172,7 @@ Init_socket(void)
 
     rb_define_method(rb_cSocket, "initialize", sock_initialize, -1);
     rb_define_method(rb_cSocket, "connect", sock_connect, 1);
-    rb_define_method(rb_cSocket, "connect_nonblock", sock_connect_nonblock, 1);
+    rb_define_method(rb_cSocket, "connect_nonblock", sock_connect_nonblock, -1);
     rb_define_method(rb_cSocket, "bind", sock_bind, 1);
     rb_define_method(rb_cSocket, "listen", rsock_sock_listen, 1);
     rb_define_method(rb_cSocket, "accept", sock_accept, 0);
@@ -2187,4 +2201,8 @@ Init_socket(void)
 #endif
 
     rb_define_singleton_method(rb_cSocket, "ip_address_list", socket_s_ip_address_list, 0);
+
+#undef rb_intern
+    sym_exception = ID2SYM(rb_intern("exception"));
+    sym_wait_writable = ID2SYM(rb_intern("wait_writable"));
 }
diff --git a/test/openssl/test_pair.rb b/test/openssl/test_pair.rb
index 2c2776b..cfcfa7a 100644
--- a/test/openssl/test_pair.rb
+++ b/test/openssl/test_pair.rb
@@ -283,7 +283,7 @@ module OpenSSL::TestPairM
     serv.close if serv && !serv.closed?
   end
 
-  def test_accept_nonblock_no_exception
+  def test_connect_accept_nonblock_no_exception
     ctx2 = OpenSSL::SSL::SSLContext.new
     ctx2.ciphers = "ADH"
     ctx2.tmp_dh_callback = proc { OpenSSL::TestUtils::TEST_KEY_DH1024 }
@@ -297,11 +297,31 @@ module OpenSSL::TestPairM
     ctx1 = OpenSSL::SSL::SSLContext.new
     ctx1.ciphers = "ADH"
     s1 = OpenSSL::SSL::SSLSocket.new(sock1, ctx1)
-    th = Thread.new { s1.connect }
+    th = Thread.new do
+      rets = []
+      begin
+        rv = s1.connect_nonblock(exception: false)
+        rets << rv
+        case rv
+        when :wait_writable
+          IO.select(nil, [s1], nil, 5)
+        when :wait_readable
+          IO.select([s1], nil, nil, 5)
+        end
+      end until rv == s1
+      rets
+    end
+
     until th.join(0.01)
       accepted = s2.accept_nonblock(exception: false)
       assert_includes([s2, :wait_readable, :wait_writable ], accepted)
     end
+
+    rets = th.value
+    assert_instance_of Array, rets
+    rets.each do |rv|
+      assert_includes([s1, :wait_readable, :wait_writable ], rv)
+    end
   ensure
     s1.close if s1
     s2.close if s2
diff --git a/test/socket/test_nonblock.rb b/test/socket/test_nonblock.rb
index 94ed198..6912046 100644
--- a/test/socket/test_nonblock.rb
+++ b/test/socket/test_nonblock.rb
@@ -55,6 +55,29 @@ class TestSocketNonblock < Test::Unit::TestCase
     s.close if s
   end
 
+  def test_connect_nonblock_no_exception
+    serv = Socket.new(:INET, :STREAM)
+    serv.bind(Socket.sockaddr_in(0, "127.0.0.1"))
+    serv.listen(5)
+    c = Socket.new(:INET, :STREAM)
+    servaddr = serv.getsockname
+    rv = c.connect_nonblock(servaddr, exception: false)
+    case rv
+    when 0
+      # some OSes return immediately on non-blocking local connect()
+    else
+      assert_equal :wait_writable, rv
+    end
+    assert_equal([ [], [c], [] ], IO.select(nil, [c], nil, 60))
+    s, sockaddr = serv.accept
+    assert_equal(Socket.unpack_sockaddr_in(c.getsockname),
+                 Socket.unpack_sockaddr_in(sockaddr))
+  ensure
+    serv.close if serv
+    c.close if c
+    s.close if s
+  end
+
   def test_udp_recvfrom_nonblock
     u1 = UDPSocket.new
     u2 = UDPSocket.new
-- 
EW