Better error message on connection closed

lib/bloomrb.rb

@@ -1,6 +1,8 @@
 require 'socket'
 
 class Bloomrb
+  class ConnectionClosed < Exception; end
+
   attr_accessor :host, :port, :retries
 
   def initialize(host = 'localhost', port = 8673, retries = 5)
@@ -97,27 +99,34 @@ def execute *args
 
     cmd = args.join(' ')
     cmd += ' ' + opts.map{|k, v| "#{k}=#{v}"}.join(' ') unless opts.empty?
-    
+
     retry_count = 0
     begin
       socket.puts(cmd)
-      result = socket.gets.chomp
+      result = read_socket_data
       raise "#{result}: #{cmd[0..99]}" if result =~ /^Client Error:/
 
       if result == 'START'
         result = []
-        while (s = socket.gets.chomp) != 'END'
+        while (s = read_socket_data) != 'END'
           result << s
         end
       end
       result
-    rescue Errno::ECONNRESET, Errno::ECONNABORTED, Errno::ECONNREFUSED, Errno::EPIPE
+    rescue Errno::ECONNRESET, Errno::ECONNABORTED, Errno::ECONNREFUSED, Errno::EPIPE, ConnectionClosed
       raise if (retry_count += 1) >= retries
       @socket = nil
       sleep(1)
       retry
     end
   end
+
+  def read_socket_data
+    received_data = socket.gets
+    # IO#gets returns nil on EOF
+    raise ConnectionClosed.new("Connection has been closed") if received_data.nil?
+    received_data.tap {|rd| rd.chomp! }
+  end
 end
 
 class BloomFilter

test/test_bloomrb.rb

@@ -162,14 +162,22 @@ class BloomrbTest < Test::Unit::TestCase
         'storage' => '1797211'}, @bloom.info('foobar'))
     end
 
-    should "retry" do
+    should "retry on connection error" do
       @bloom.expects(:sleep).with(1)
       @socket.expects(:puts).twice.with("s foobar fookey")
       @socket.expects(:gets).twice.raises(Errno::ECONNRESET).then.returns("No")
 
       assert_equal false, @bloom.set('foobar', :fookey)
     end
 
+    should "retry on connection closed" do
+      @bloom.expects(:sleep).with(1)
+      @socket.expects(:puts).twice.with("s foobar fookey")
+      @socket.expects(:gets).twice.returns(nil).then.returns("No")
+
+      assert_equal false, @bloom.set('foobar', :fookey)
+    end
+
     should "not retry if retries set to 0" do
       @bloom.retries = 0
 
@@ -181,7 +189,7 @@ class BloomrbTest < Test::Unit::TestCase
       end
     end
 
-    should "raise after 5 retries" do
+    should "raise after 5 retries on connection error" do
       @bloom.expects(:sleep).times(4).with(1)
       @socket.expects(:puts).times(5).with("s foobar fookey")
       @socket.expects(:gets).times(5).raises(Errno::ECONNRESET)
@@ -190,5 +198,15 @@ class BloomrbTest < Test::Unit::TestCase
         @bloom.set('foobar', :fookey)
       end
     end
+
+    should "raise after 5 retries on connection closed" do
+      @bloom.expects(:sleep).times(4).with(1)
+      @socket.expects(:puts).times(5).with("s foobar fookey")
+      @socket.expects(:gets).times(5).returns(nil)
+
+      assert_raises Bloomrb::ConnectionClosed do
+        @bloom.set('foobar', :fookey)
+      end
+    end
   end
 end