network: make verify_uris method more robust

IO.select can raise errors (IOError most notably) on bad
descriptors.  Ensure we can detect and delete them from our
waiting sets if IO.select raises something.

2 files changed, 100 insertions, 32 deletions

diff --git a/lib/mogilefs/network.rb b/lib/mogilefs/network.rb
index 7618c26..0350749 100644
--- a/lib/mogilefs/network.rb
+++ b/lib/mogilefs/network.rb
@@ -7,8 +7,6 @@ module MogileFS::Network
   # with the expected HTTP code within the timeout period (in seconds).
   def verify_uris(uris = [], expect = '200', timeout = 2.00)
     uri_socks = {}
-    ok_uris = []
-    sockets = []
 
     # first, we asynchronously connect to all of them
     uris.each do |uri|
@@ -18,47 +16,89 @@ module MogileFS::Network
 
     # wait for at least one of them to finish connecting and send
     # HTTP requests to the connected ones
-    begin
-      t0 = Time.now
-      r = IO.select(nil, uri_socks.keys, nil, timeout > 0 ? timeout : 0)
-      timeout -= (Time.now - t0)
-      break unless r && r[1]
-      r[1].each do |sock|
-        begin
-          # we don't about short/interrupted writes here, if the following
-          # request fails or blocks then the server is flat-out hopeless
-          sock.syswrite "HEAD #{uri_socks[sock].request_uri} HTTP/1.0\r\n\r\n"
-          sockets << sock
-        rescue
-          sock.close rescue nil
-        end
-      end
-    end until sockets[0] || timeout < 0
+    sockets, timeout = get_writable_set(uri_socks, timeout)
 
     # Await a response from the sockets we had written to, we only need one
     # valid response, but we'll take more if they return simultaneously
-    if sockets[0]
+    sockets[0] ? get_readable_uris(sockets, uri_socks, expect, timeout) : []
+
+    ensure
+      uri_socks.keys.each { |sock| sock.close rescue nil }
+  end
+
+  private
+    include MogileFS::Util
+
+    # returns an array of writeable Sockets and leftover from timeout
+    def get_writable_set(uri_socks, timeout)
+      sockets = []
       begin
         t0 = Time.now
-        r = IO.select(sockets, nil, nil, timeout > 0 ? timeout : 0)
+        r = begin
+         IO.select(nil, uri_socks.keys, nil, timeout > 0 ? timeout : 0)
+        rescue
+          # get rid of bad descriptors
+          uri_socks.delete_if do |sock, uri|
+            begin
+              sock.recv_nonblock(1)
+              false # should never get here for HTTP, really...
+            rescue Errno::EAGAIN, Errno::EINTR
+              false
+            rescue
+              sock.close rescue nil
+              true
+            end
+          end
+          timeout -= (Time.now - t0)
+          retry if timeout >= 0
+        end
+
+        break unless r && r[1]
+
+        r[1].each do |sock|
+          begin
+            # we don't about short/interrupted writes here, if the following
+            # request fails or blocks then the server is flat-out hopeless
+            sock.syswrite "HEAD #{uri_socks[sock].request_uri} HTTP/1.0\r\n\r\n"
+            sockets << sock
+          rescue
+            sock.close rescue nil
+          end
+        end
+
         timeout -= (Time.now - t0)
-        break unless r && r[0]
-        r[0].each do |sock|
-          buf = sock.recv_nonblock(128, Socket::MSG_PEEK) rescue next
+      end until (sockets[0] || timeout < 0)
+
+      [ sockets, timeout ]
+    end
+
+    # returns an array of URIs from uri_socks that are good
+    def get_readable_uris(sockets, uri_socks, expect, timeout)
+      ok_uris = []
+
+      begin
+        t0 = Time.now
+        r = IO.select(sockets, nil, nil, timeout > 0 ? timeout : 0) rescue nil
+
+        (r && r[0] ? r[0] : sockets).each do |sock|
+          buf = begin
+            sock.recv_nonblock(128, Socket::MSG_PEEK)
+          rescue Errno::EAGAIN, Errno::EINTR
+            next
+          rescue
+            sockets.delete(sock) # socket went bad
+            next
+          end
+
           if buf && /\AHTTP\/[\d\.]+ #{expect} / =~ buf
             ok_uris << uri_socks.delete(sock)
             sock.close rescue nil
           end
         end
-      end
-    end until ok_uris[0] || timeout < 0
-
-    ok_uris
-    ensure
-      uri_socks.keys.each { |sock| sock.close rescue nil }
-  end
+        timeout -= (Time.now - t0)
+      end until ok_uris[0] || timeout < 0
 
-  private
-    include MogileFS::Util
+      ok_uris
+    end
 
 end # module MogileFS::Network
diff --git a/test/test_network.rb b/test/test_network.rb
index f479ca0..94b51fe 100644
--- a/test/test_network.rb
+++ b/test/test_network.rb
@@ -24,4 +24,32 @@ class TestNetwork < Test::Unit::TestCase
       TempServer.destroy_all!
   end
 
+  def test_verify_non_existent
+    good = TempServer.new(Proc.new do |serv,port|
+      client,client_addr = serv.accept
+      client.readpartial(4096)
+      client.syswrite("HTTP/1.0 200 OK\r\nContent-Length: 0\r\n\r\n")
+    end)
+    bad = TempServer.new(Proc.new { |serv,port| serv.close })
+
+    good_uri = URI.parse("http://127.0.0.1:#{good.port}/")
+    bad_uri = URI.parse("http://127.0.0.1:#{bad.port}/")
+    ok = verify_uris([ good_uri, bad_uri ])
+    assert_equal [ good_uri ], ok
+    ensure
+      TempServer.destroy_all!
+  end
+
+  def test_verify_all_bad
+    good = TempServer.new(Proc.new { |serv,port| serv.close })
+    bad = TempServer.new(Proc.new { |serv,port| serv.close })
+
+    good_uri = URI.parse("http://127.0.0.1:#{good.port}/")
+    bad_uri = URI.parse("http://127.0.0.1:#{bad.port}/")
+    ok = verify_uris([ good_uri, bad_uri ], '200', 1.0)
+    assert ok.empty?, "nothing returned"
+    ensure
+      TempServer.destroy_all!
+  end
+
 end