From 2c45407690ed84c1266b551f0779eab20a5e37e9 Mon Sep 17 00:00:00 2001
From: Eric Wong <e@80x24.org>
Date: Fri, 24 Mar 2017 05:46:46 +0000
Subject: [PATCH] socket: avoid fcntl for read/write_nonblock on Linux

On platforms where MSG_DONTWAIT works reliably on all sockets
(so far, I know of Linux), we can avoid fcntl syscalls and
implement IO#write_nonblock and IO#read_nonblock in terms of the
socket-specific send and recv family of syscalls.

This avoids side effects on the socket, and also encourages
generic code to be written in cases where IO wrappers like
OpenSSL::SSL::SSLSocket are used.
---
 ext/socket/lib/socket.rb        | 18 ++++++++++++++++++
 test/socket/test_basicsocket.rb | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/ext/socket/lib/socket.rb b/ext/socket/lib/socket.rb
index ad7d1e7aa0..ff20453207 100644
--- a/ext/socket/lib/socket.rb
+++ b/ext/socket/lib/socket.rb
@@ -442,6 +442,24 @@ def recvmsg_nonblock(dlen = nil, flags = 0, clen = nil,
                        scm_rights: false, exception: true)
     __recvmsg_nonblock(dlen, flags, clen, scm_rights, exception)
   end
+
+  # Linux-specific optimizations to avoid fcntl for IO#read_nonblock
+  # and IO#write_nonblock using MSG_DONTWAIT
+  # Do other platforms suport MSG_DONTWAIT reliably?
+  if RUBY_PLATFORM =~ /linux/ && Socket.const_defined?(:MSG_DONTWAIT)
+    def read_nonblock(len, str = nil, exception: true) # :nodoc:
+      case rv = __recv_nonblock(len, 0, str, exception)
+      when '' # recv_nonblock returns empty string on EOF
+        exception ? raise(EOFError, 'end of file reached') : nil
+      else
+        rv
+      end
+    end
+
+    def write_nonblock(buf, exception: true) # :nodoc:
+      __sendmsg_nonblock(buf, 0, nil, nil, exception)
+    end
+  end
 end
 
 class Socket < BasicSocket
diff --git a/test/socket/test_basicsocket.rb b/test/socket/test_basicsocket.rb
index e17a675d8a..0bd0408d83 100644
--- a/test/socket/test_basicsocket.rb
+++ b/test/socket/test_basicsocket.rb
@@ -152,4 +152,40 @@ def test_for_fd
       sock.close
     end
   end
+
+  def test_read_write_nonblock
+    socks do |sserv, ssock, csock|
+      buf = String.new
+      assert_equal :wait_readable, ssock.read_nonblock(1, buf, exception: false)
+      assert_equal 5, csock.write_nonblock('hello')
+      IO.select([ssock])
+      assert_same buf, ssock.read_nonblock(5, buf, exception: false)
+      assert_equal 'hello', buf
+      buf = '*' * 16384
+      n = 0
+
+      case w = csock.write_nonblock(buf, exception: false)
+      when Integer
+        n += w
+      when :wait_writable
+        break
+      end while true
+
+      assert_equal :wait_writable, w
+      assert_raise(IO::WaitWritable) { loop { csock.write_nonblock(buf) } }
+      assert_operator n, :>, 0
+      csock.close
+
+      case r = ssock.read_nonblock(16384, buf, exception: false)
+      when String
+        next
+      when nil
+        break
+      else
+        flunk "unexpected read_nonblock return: #{r.inspect}"
+      end while true
+
+      assert_raise(EOFError) { ssock.read_nonblock(1) }
+    end
+  end
 end if defined?(BasicSocket)
-- 
EW