aggregate/last_data_recv: support Socket#accept{,_nonblock}

Socket#accept and Socket#accept_nonblock return an Addrinfo
object in addition to a client socket.  This allows web servers
to avoid having to make getpeername(2) syscalls to get the
same information.

2 files changed, 60 insertions, 6 deletions

diff --git a/lib/raindrops/aggregate/last_data_recv.rb b/lib/raindrops/aggregate/last_data_recv.rb
index 6919fbc..32908f2 100644
--- a/lib/raindrops/aggregate/last_data_recv.rb
+++ b/lib/raindrops/aggregate/last_data_recv.rb
@@ -10,6 +10,8 @@ require "socket"
 # Methods wrapped include:
 # - TCPServer#accept
 # - TCPServer#accept_nonblock
+# - Socket#accept
+# - Socket#accept_nonblock
 # - Kgio::TCPServer#kgio_accept
 # - Kgio::TCPServer#kgio_tryaccept
 module Raindrops::Aggregate::LastDataRecv
@@ -33,8 +35,10 @@ module Raindrops::Aggregate::LastDataRecv
 
   # automatically extends any TCPServer objects used by Unicorn
   def self.cornify!
-    Unicorn::HttpServer::LISTENERS.each do |sock|
-      sock.extend(self) if TCPServer === sock
+    Unicorn::HttpServer::LISTENERS.each do |s|
+      if TCPServer === s || (s.instance_of?(Socket) && s.local_address.ip?)
+        s.extend(self)
+      end
     end
   end
 
@@ -60,8 +64,8 @@ module Raindrops::Aggregate::LastDataRecv
     count! super
   end
 
-  def accept_nonblock
-    count! super
+  def accept_nonblock(exception: true)
+    count! super(exception: exception)
   end
 
   # :startdoc:
@@ -72,12 +76,19 @@ module Raindrops::Aggregate::LastDataRecv
   #
   # We require TCP_DEFER_ACCEPT on the listen socket for
   # +last_data_recv+ to be accurate
-  def count!(io)
+  def count!(ret)
+    case ret
+    when :wait_readable
+    when Array # Socket#accept_nonblock
+      io = ret[0]
+    else # TCPSocket#accept_nonblock
+      io = ret
+    end
     if io
       x = Raindrops::TCP_Info.new(io)
       @raindrops_aggregate << x.last_data_recv
     end
-    io
+    ret
   end
 end
 
diff --git a/test/test_last_data_recv.rb b/test/test_last_data_recv.rb
new file mode 100644
index 0000000..ef84e05
--- /dev/null
+++ b/test/test_last_data_recv.rb
@@ -0,0 +1,43 @@
+require 'test/unit'
+require 'raindrops'
+require 'io/wait'
+
+class TestLastDataRecv < Test::Unit::TestCase
+  def test_accept_nonblock_agg
+    s = Socket.new(:INET, :STREAM, 0)
+    s.listen(128)
+    addr = s.connect_address
+    s.extend(Raindrops::Aggregate::LastDataRecv)
+    s.raindrops_aggregate = []
+    c = Socket.new(:INET, :STREAM, 0)
+    c.connect(addr)
+    c.write '.' # for TCP_DEFER_ACCEPT
+    client, ai = s.accept_nonblock(exception: false)
+    assert client.kind_of?(Socket)
+    assert ai.kind_of?(Addrinfo)
+    assert_equal 1, s.raindrops_aggregate.size
+    assert s.raindrops_aggregate[0].instance_of?(Integer)
+    client, ai = s.accept_nonblock(exception: false)
+    assert_equal :wait_readable, client
+    assert_nil ai
+    assert_equal 1, s.raindrops_aggregate.size
+    assert_raise(IO::WaitReadable) { s.accept_nonblock }
+  end
+
+  def test_accept_nonblock_one
+    s = TCPServer.new('127.0.0.1', 0)
+    s.extend(Raindrops::Aggregate::LastDataRecv)
+    s.raindrops_aggregate = []
+    addr = s.addr
+    c = TCPSocket.new(addr[3], addr[1])
+    c.write '.' # for TCP_DEFER_ACCEPT
+    client = s.accept_nonblock(exception: false)
+    assert client.kind_of?(TCPSocket)
+    assert_equal 1, s.raindrops_aggregate.size
+    assert s.raindrops_aggregate[0].instance_of?(Integer)
+    client = s.accept_nonblock(exception: false)
+    assert_equal :wait_readable, client
+    assert_equal 1, s.raindrops_aggregate.size
+    assert_raise(IO::WaitReadable) { s.accept_nonblock }
+  end
+end if RUBY_PLATFORM =~ /linux/