lib/aerospike/socket/base.rb in aerospike-2.7.0 vs lib/aerospike/socket/base.rb in aerospike-2.8.0

@@ -33,20 +33,12 @@
           bytes_read += result.bytesize
         end
       end
 
       def read_from_socket(length)
-        begin
+        with_timeout(@timeout) do
           read_nonblock(length)
-        rescue ::IO::WaitReadable => e
-          if ::IO::select([self], nil, nil, @timeout)
-            retry
-          else
-            raise ::Aerospike::Exceptions::Connection.new("#{e}")
-          end
-        rescue => e
-          raise ::Aerospike::Exceptions::Connection.new("#{e}")
         end
       end
 
       def write(buffer, length)
         bytes_written = 0
@@ -54,20 +46,12 @@
           bytes_written += write_to_socket(buffer.read(bytes_written, length - bytes_written))
         end
       end
 
       def write_to_socket(data)
-        begin
+        with_timeout(@timeout) do
           write_nonblock(data)
-        rescue ::IO::WaitWritable => e
-          if ::IO::select(nil, [self], nil, @timeout)
-            retry
-          else
-            raise ::Aerospike::Exceptions::Connection.new("#{e}")
-          end
-        rescue => e
-          raise ::Aerospike::Exceptions::Connection.new("#{e}")
         end
       end
 
       def timeout=(timeout)
         @timeout = timeout && timeout > 0 ? timeout : nil
@@ -75,12 +59,55 @@
 
       def connected?
         !closed?
       end
 
+      # Returns whether the connection to the server is alive.
+      #
+      # It is useful to call this method before making a call to the server
+      # that would change data on the server.
+      #
+      # Note: This method is only useful if the server closed the connection or
+      # if a previous connection failure occurred. If the server is hard killed
+      # this will still return true until one or more writes are attempted.
+      def alive?
+        return false if closed?
+
+        if IO.select([self], nil, nil, 0)
+          !eof? rescue false
+        else
+          true
+        end
+      rescue IOError
+        false
+      end
+
       def close
         return if closed?
         super()
+      end
+
+      private
+
+      # Note: For SSL connections, read_nonblock may invoke write system call,
+      # which may raise IO::WaitWritable, and vice versa, due to SSL
+      # renegotiation, so we should always rescue both.
+      def with_timeout(timeout, &block)
+        block.call
+      rescue IO::WaitReadable => e
+        if IO::select([self], nil, nil, timeout)
+          retry
+        else
+          fail Aerospike::Exceptions::Connection, "Socket timeout: #{e}"
+        end
+      rescue IO::WaitWritable => e
+        if IO::select(nil, [self], nil, timeout)
+          retry
+        else
+          fail Aerospike::Exceptions::Connection, "Socket timeout: #{e}"
+        end
+      rescue => e
+        raise Aerospike::Exceptions::Connection, "Socket error: #{e}"
       end
     end
   end
 end
\ No newline at end of file