agent.rb in ftw-0.0.7

- old
+ new

@@ -1,8 +1,9 @@
 require "ftw/namespace"
 require "ftw/request"
 require "ftw/connection"
+require "ftw/protocol"
 require "ftw/pool"
 require "ftw/websocket"
 require "addressable/uri"
 require "cabin"
 require "logger"
@@ -33,10 +34,13 @@
 #     response = agent.get!("http://www.google.com/")
 #     puts response.body.head
 #
 # TODO(sissel): TBD: implement cookies... delicious chocolate chip cookies.
 class FTW::Agent
+  include FTW::Protocol
+
+  # List of standard HTTP methods described in RFC2616
   STANDARD_METHODS = %w(options get head post put delete trace connect)
 
   # Everything is private by default.
   # At the bottom of this class, public methods will be declared.
   private
@@ -93,29 +97,20 @@
   # Make a new websocket connection.
   #
   # This will send the http request. If the websocket handshake
   # is successful, a FTW::WebSocket instance will be returned.
   # Otherwise, a FTW::Response will be returned.
+  #
+  # See {#request} for what the 'uri' and 'options' parameters should be.
   def websocket!(uri, options={})
     # TODO(sissel): Use FTW::Agent#upgrade! ?
     req = request("GET", uri, options)
     ws = FTW::WebSocket.new(req)
     response = execute(req)
     if ws.handshake_ok?(response)
       # response.body is a FTW::Connection
       ws.connection = response.body
-
-      # TODO(sissel): Investigate this bug
-      # There seems to be a bug in http_parser.rb (or in this library) where
-      # websocket responses lead with a newline for some reason. Work around
-      # it.
-      data = response.body.read
-      if data[0] == "\n"
-        response.body.pushback(data[1..-1])
-      else
-        response.body.pushback(data)
-      end
       return ws
     else
       return response
     end
   end # def websocket!
@@ -156,16 +151,20 @@
   # of this Request, it will be reused. Otherwise, a new connection
   # is opened.
   #
   # Redirects are always followed.
   #
-  # @params
+  # @param [FTW::Request]
   # @return [FTW::Response] the response for this request.
   def execute(request)
     # TODO(sissel): Make redirection-following optional, but default.
 
-    connection = connect(request.headers["Host"], request.port)
+    connection, error = connect(request.headers["Host"], request.port)
+    if !error.nil?
+      p :error => error
+      raise error
+    end
     connection.secure if request.protocol == "https"
     response = request.execute(connection)
 
     redirects = 0
     while response.redirect? and response.headers.include?("Location")
@@ -190,11 +189,16 @@
       # agent's cookie store
 
       @logger.debug("Redirecting", :location => response.headers["Location"])
       redirects += 1
       request.use_uri(response.headers["Location"])
-      connection = connect(request.headers["Host"], request.port)
+      connection, error = connect(request.headers["Host"], request.port)
+      # TODO(sissel): Do better error handling than raising.
+      if !error.nil?
+        p :error => error
+        raise error
+      end
       connection.secure if request.protocol == "https"
       response = request.execute(connection)
     end
 
     # RFC 2616 section 9.4, HEAD requests MUST NOT have a message body.
@@ -211,18 +215,31 @@
 
   # Returns a FTW::Connection connected to this host:port.
   def connect(host, port)
     address = "#{host}:#{port}"
     @logger.debug("Fetching from pool", :address => address)
+    error = nil
     connection = @pool.fetch(address) do
       @logger.info("New connection to #{address}")
       connection = FTW::Connection.new(address)
-      connection.connect
-      connection
+      error = connection.connect
+      if !error.nil?
+        # Return nil to the pool, so like, we failed..
+        nil
+      else
+        # Otherwise return our new connection
+        connection
+      end
     end
+
+    if !error.nil?
+      @logger.error("Connection failed", :destination => address, :error => error)
+      return nil, error
+    end
+
     @logger.debug("Pool fetched a connection", :connection => connection)
     connection.mark
-    return connection
+    return connection, nil
   end # def connect
 
   public(:initialize, :execute, :websocket!, :upgrade!)
 end # class FTW::Agent