From 766e58508069a6ef291d4a49fdfa567ccd2bb61f Mon Sep 17 00:00:00 2001
From: kazjote <kazjote@gmail.com>
Date: Wed, 23 Jan 2013 15:57:06 +0100
Subject: [PATCH] Better error message on connection closed

---
 .gitignore           |  1 +
 lib/bloomrb.rb       | 17 +++++++++++++----
 test/test_bloomrb.rb | 22 ++++++++++++++++++++--
 3 files changed, 34 insertions(+), 6 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b844b14
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+Gemfile.lock
diff --git a/lib/bloomrb.rb b/lib/bloomrb.rb
index 1507fe4..39b3cb6 100644
--- a/lib/bloomrb.rb
+++ b/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)
@@ -92,27 +94,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
       throw result 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
diff --git a/test/test_bloomrb.rb b/test/test_bloomrb.rb
index 64f593f..4425989 100644
--- a/test/test_bloomrb.rb
+++ b/test/test_bloomrb.rb
@@ -162,7 +162,7 @@ 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")
@@ -170,6 +170,14 @@ class BloomrbTest < Test::Unit::TestCase
       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