connection.rb in beanstalk-client-0.6.1

- old
+ new

@@ -17,10 +17,11 @@
 
 require 'socket'
 require 'fcntl'
 require 'yaml'
 require 'beanstalk-client/bag'
+require 'beanstalk-client/errors'
 require 'beanstalk-client/job'
 
 module Beanstalk
   class RawConnection
     attr_reader :addr
@@ -181,11 +182,11 @@
     end
 
     def method_missing(selector, *args, &block)
       begin
         @conn.send(selector, *args, &block)
-      rescue EOFError, Errno::ECONNRESET => ex
+      rescue EOFError, Errno::ECONNRESET, Errno::EPIPE, UnexpectedResponse => ex
         @multi.remove(@conn)
         raise ex
       end
     end
   end
@@ -214,29 +215,30 @@
 
     def open_connections()
       @connections.values()
     end
 
-    def put(body, pri=65536, delay=0, ttr=120)
-      pick_connection.put(body, pri, delay, ttr)
-    rescue EOFError, Errno::ECONNRESET
+    def send_to_conn(sel, *args)
+      pick_connection.send(sel, *args)
+    rescue DrainingError
+      # Don't reconnect -- we're not interested in this server
+      retry
+    rescue EOFError, Errno::ECONNRESET, Errno::EPIPE
       connect()
       retry
     end
 
+    def put(body, pri=65536, delay=0, ttr=120)
+      send_to_conn(:put, body, pri, delay, ttr)
+    end
+
     def yput(obj, pri=65536, delay=0, ttr=120)
-      pick_connection.yput(obj, pri, delay, ttr)
-    rescue EOFError, Errno::ECONNRESET
-      connect()
-      retry
+      send_to_conn(:yput, obj, pri, delay, ttr)
     end
 
     def reserve()
-      pick_connection.reserve()
-    rescue EOFError, Errno::ECONNRESET
-      connect()
-      retry
+      send_to_conn(:reserve)
     end
 
     def raw_stats()
       Hash[*@connections.map{|a,c| [a, c.stats()]}.inject([]){|a,b|a+b}]
     end
@@ -258,9 +260,9 @@
     end
 
     private
 
     def pick_connection()
-      open_connections[rand(open_connections.size)] or raise 'not connected'
+      open_connections[rand(open_connections.size)] or raise NotConnected
     end
   end
 end