lib/kjess/socket.rb in kjess-1.1.0 vs lib/kjess/socket.rb in kjess-1.2.0

@@ -1,11 +1,11 @@
 require 'socket'
 
 module KJess
   # Socket: A specialized socket that has been configure
   class Socket
-    class Error < KJess::Error; end
+    class Error < KJess::NetworkError; end
     class Timeout < Error; end
 
     # Internal:
     # The timeout for reading in seconds. Defaults to 2
     attr_accessor :read_timeout
@@ -19,11 +19,11 @@
     attr_reader :write_timeout
 
     # Internal:
     # The host this socket is connected to
     attr_reader :host
-    
+
     # Internal:
     # The port this socket is connected to
     attr_reader :port
 
     # Internal
@@ -96,11 +96,11 @@
     def keepalive_active?
       @keepalive_active
     end
 
     # Internal: Low level socket allocation and option configuration
-    # 
+    #
     # Using the options from the initializer, a new ::Socket is created that
     # is:
     #
     #   TCP, IPv4 only, autoclosing on exit, nagle's algorithm is disabled and has
     #   TCP Keepalive options set if keepalive is supported.
@@ -165,16 +165,16 @@
       addrs      = ::Socket.getaddrinfo(host, port, ::Socket::AF_INET, ::Socket::SOCK_STREAM )
       errors     = []
       conn_error = lambda { raise errors.first }
       sock       = nil
 
-      addrs.find( conn_error ) do |addr| 
+      addrs.find( conn_error ) do |addr|
         sock = connect_or_error( addr, deadline, errors )
       end
       return sock
     end
-  
+
     # Internal: Connect to the destination or raise an error.
     #
     # Connect to the address or capture the error of the connection
     #
     # addr     - An address returned from Socket.getaddrinfo()
@@ -206,26 +206,27 @@
       sockaddr = ::Socket.pack_sockaddr_in(addr[1], addr[3])
       sock     = blank_socket()
       sock.connect_nonblock( sockaddr )
       return sock
     rescue Errno::EINPROGRESS
-      if IO.select(nil, [sock], nil, timeout).nil? then
+      if !wait_writable(timeout, sock)
         raise Timeout, "Could not connect to #{host}:#{port} within #{timeout} seconds"
       end
       return connect_nonblock_finalize( sock, sockaddr )
     rescue => ex
       raise Error, "Could not connect to #{host}:#{port}: #{ex.class}: #{ex.message}", ex.backtrace
     end
 
 
     # Internal: Make sure that a non-blocking connect has truely connected.
-    # 
+    #
     # Ensure that the given socket is actually connected to the given adddress.
     #
     # Returning the socket if it is and raising an Error if it isn't.
     def connect_nonblock_finalize( sock, sockaddr )
       sock.connect_nonblock( sockaddr )
+      return sock
     rescue Errno::EISCONN
       return sock
     rescue => ex
       raise Error, "Could not connect to #{host}:#{port}: #{ex.class}: #{ex.message}", ex.backtrace
     end
@@ -258,11 +259,11 @@
     #
     # Returns the bytes read
     def readpartial(maxlen, outbuf = nil)
       return socket.read_nonblock(maxlen, outbuf)
     rescue Errno::EWOULDBLOCK, Errno::EAGAIN, Errno::ECONNRESET
-      if IO.select([@socket], nil, nil, read_timeout)
+      if wait_readable(read_timeout)
         retry
       else
         raise Timeout, "Could not read from #{host}:#{port} in #{read_timeout} seconds"
       end
     end
@@ -274,18 +275,26 @@
     # Raises an error if it is unable to write the data to the socket within the
     # write_timeout.
     #
     # returns nothing
     def write( buf )
-      until buf.length == 0
+      until buf.nil? or (buf.length == 0) do
         written = socket.write_nonblock(buf)
         buf = buf[written, buf.length]
       end
     rescue Errno::EWOULDBLOCK, Errno::EINTR, Errno::EAGAIN, Errno::ECONNRESET
-      if IO.select(nil, [socket], nil, write_timeout)
+      if wait_writable(write_timeout)
         retry
       else
         raise Timeout, "Could not write to #{host}:#{port} in #{write_timeout} seconds"
       end
+    end
+
+    def wait_writable(timeout = nil, socket = @socket)
+      IO.select(nil, [socket], nil, timeout || write_timeout)
+    end
+
+    def wait_readable(timeout = nil, socket = @socket)
+      IO.select([socket], nil, nil, timeout || read_timeout)
     end
   end
 end