kronk 1.6.2 → 1.7.0

data/lib/kronk/player.rb

@@ -6,33 +6,57 @@ class Kronk
   # Kronk includes a Suite (test-like) output, a Stream (chunked) output,
   # and a Benchmark output.
 
-  class Player
+  class Player < QueueRunner
 
-    attr_accessor :number, :concurrency, :queue, :count, :input,
-                  :output, :mutex, :threads, :reader_thread
+    RESCUABLE = [Kronk::Exception, Errno::ECONNRESET, URI::InvalidURIError]
+
+    attr_accessor :input, :output
 
     ##
     # Create a new Player for batch diff or response validation.
     # Supported options are:
-    # :concurrency:: Fixnum - The maximum number of concurrent requests to make
-    # :number:: Fixnum - The number of requests to make
     # :io:: IO - The IO instance to read from
     # :output:: Class - The output class to use (see Player::Output)
     # :parser:: Class - The IO parser to use.
 
     def initialize opts={}
-      @number      = opts[:number]
-      @concurrency = opts[:concurrency]
-      @concurrency = 1 if !@concurrency || @concurrency <= 0
+      super
       self.output_from opts[:output] || Suite
 
-      @count     = 0
-      @queue     = []
-      @threads   = []
-      @input     = InputReader.new opts[:io], opts[:parser]
-      @last_req  = nil
+      @input      = InputReader.new opts[:io], opts[:parser]
+      @use_output = true
+      @last_req   = nil
+
+      @on_input   = Proc.new do
+        stop_input! if !@number && @input.eof?
+        @last_req = @input.get_next || @queue.last || @last_req || {}
+      end
+
+      @on_result  = Proc.new do |kronk, err, mutex|
+        err ? @output.error(err, kronk, mutex) :
+              @output.result(kronk, mutex)
+      end
 
-      @mutex = Mutex.new
+      on(:input, &@on_input)
+      on(:interrupt){
+        @output.completed if @use_output
+        exit 2
+      }
+      on(:start){ @output.start if @use_output }
+      on(:complete){ @output.completed if @use_output }
+    end
+
+
+    ##
+    # Populate the queue by reading from the given IO instance and
+    # parsing it into kronk options.
+    #
+    # Default parser is RequestParser. See InputReader for parser requirements.
+
+    def from_io io, parser=nil
+      @input.io     = io
+      @input.parser = parser if parser
+      @input
     end
 
 
@@ -56,37 +80,17 @@ class Kronk
     end
 
 
-    ##
-    # Adds kronk request hash options to queue.
-    # See Kronk#compare for supported options.
-
-    def queue_req kronk_opts
-      @queue << kronk_opts
-    end
-
-
-    ##
-    # Populate the queue by reading from the given IO instance and
-    # parsing it into kronk options.
-    #
-    # Default parser is RequestParser. See InputReader for parser requirements.
-
-    def from_io io, parser=nil
-      @input.io     = io
-      @input.parser = parser if parser
-      @input
-    end
-
-
     ##
     # Process the queue to compare two uris.
     # If options are given, they are merged into every request.
 
-    def compare uri1, uri2, opts={}
+    def compare uri1, uri2, opts={}, &block
       return Cmd.compare uri1, uri2, @queue.shift.merge(opts) if single_request?
 
-      process_queue do |kronk_opts|
-        process_compare uri1, uri2, kronk_opts.merge(opts)
+      method = self.class.async ? :process_one_async : :process_one
+
+      run do |kronk_opts, mutex|
+        send method, kronk_opts.merge(opts), 'compare', uri1, uri2, &block
       end
     end
 
@@ -95,139 +99,85 @@ class Kronk
     # Process the queue to request uris.
     # If options are given, they are merged into every request.
 
-    def request uri, opts={}
-      return Cmd.request(uri, @queue.shift.merge(opts)) if single_request?
+    def request uri, opts={}, &block
+      return Cmd.request uri, @queue.shift.merge(opts) if single_request?
 
-      process_queue do |kronk_opts|
-        process_request uri, kronk_opts.merge(opts)
-      end
-    end
+      method = self.class.async ? :process_one_async : :process_one
 
-
-    ##
-    # Check if we're only processing a single case.
-    # If so, yield a single item and return immediately.
-
-    def single_request?
-      @queue << next_request if @queue.empty? && (!@number || @number <= 1)
-      @queue.length == 1 && @input.eof?
+      run do |kronk_opts, mutex|
+        send method, kronk_opts.merge(opts), 'request', uri, &block
+      end
     end
 
 
     ##
-    # Start processing the queue and reading from IO if available.
-    # Calls Output#start method and returns the value of Output#completed
-    # once processing is finished.
+    # Run a single compare or request and call the Output#result or
+    # Output#error method.
     #
-    # Yields queue item until queue and io (if available) are empty and the
-    # totaly number of requests to run is met (if number is set).
-
-    def process_queue
-      @reader_thread = try_fill_queue
-
-      trap 'INT' do
-        @threads.each{|t| t.kill}
-        @threads.clear
-        @reader_thread.kill
-        output_results
-        exit 2
-      end
-
-      @output.start
-      @count = 0
-
-      until finished?
-        @threads.delete_if{|t| !t.alive? }
-        next if @threads.length >= @concurrency || @queue.empty?
+    # If given a block, will yield the Kronk instance and error. If
+    # a third argument is given, mutex will also be passed and the
+    # block won't be called from a mutex lock.
+    #
+    # Returns the result of the block or of the called Output method.
 
-        @threads << Thread.new(@queue.shift) do |kronk_opts|
-          yield kronk_opts if block_given?
-        end
+    def process_one opts={}, *args, &block
+      err   = nil
+      kronk = Kronk.new opts
 
-        @count += 1
+      begin
+        kronk.send(*args)
+      rescue *RESCUABLE => e
+        err = e
       end
 
-      @threads.each{|t| t.join}
-      @threads.clear
-
-      @reader_thread.kill
-
-      output_results
+      trigger_result kronk, err, &block
     end
 
 
     ##
-    # Attempt to fill the queue by reading from the IO instance.
-    # Starts a new thread and returns the thread instance.
-
-    def try_fill_queue
-      Thread.new do
-        loop do
-          break if !@number && @input.eof?
-          next  if @queue.length >= @concurrency * 2
+    # Run a single compare or request and call the Output#result or
+    # Output#error method using EventMachine.
+    #
+    # If given a block, will yield the Kronk instance and error. If
+    # a third argument is given, mutex will also be passed and the
+    # block won't be called from a mutex lock.
+    #
+    # Returns either a EM::MultiRequest or an EM::Connection handler.
 
-          max_new = @concurrency * 2 - @queue.length
+    def process_one_async opts={}, *args, &block
+      kronk  = Kronk.new opts
+      method = args.shift.to_s + '_async'
 
-          max_new.times do
-            @queue << next_request
-            break if !@number && @input.eof?
-          end
-        end
+      kronk.send(method, *args) do |obj, err|
+        raise err if err && !RESCUABLE.find{|eclass| eclass === err}
+        trigger_result kronk, err, &block
       end
     end
 
 
     ##
-    # Gets the next request to perform and always returns a Hash.
-    # Tries from input first, then from the last item in the queue.
-    # If both fail, returns an empty Hash.
-
-    def next_request
-      @last_req = @input.get_next || @queue.last || @last_req || Hash.new
-    end
+    # Trigger a single kronk result callback.
 
+    def trigger_result kronk, err, &block
+      block ||= @on_result
 
-    ##
-    # Returns true if processing queue should be stopped, otherwise false.
-
-    def finished?
-      (@number && @count >= @number) || @queue.empty? &&
-      @input.eof? && @count > 0 && !@reader_thread.alive?
-    end
-
-
-    ##
-    # Process and output the results.
-    # Calls Output#completed method.
-
-    def output_results
-      @output.completed
-    end
-
-
-    ##
-    # Run a single compare and call the Output#result or Output#error method.
-
-    def process_compare uri1, uri2, opts={}
-      kronk = Kronk.new opts
-      kronk.compare uri1, uri2
-      @output.result kronk, @mutex
-
-    rescue Kronk::Exception, Response::MissingParser, Errno::ECONNRESET => e
-      @output.error e, kronk, @mutex
+      if block.arity > 2 || block.arity < 0
+        block.call kronk, err, @mutex
+      else
+        @mutex.synchronize do
+          block.call kronk, err
+        end
+      end
     end
 
 
     ##
-    # Run a single request and call the Output#result or Output#error method.
-
-    def process_request uri, opts={}
-      kronk = Kronk.new opts
-      kronk.retrieve uri
-      @output.result kronk, @mutex
+    # Check if w6ce5e2ce're only processing a single case.
+    # If so, yield a single item and return immediately.
 
-    rescue Kronk::Exception, Response::MissingParser, Errno::ECONNRESET => e
-      @output.error e, kronk, @mutex
+    def single_request?
+      @queue << trigger(:input) if @queue.empty? && (!@number || @number <= 1)
+      @queue.length == 1 && @triggers[:input] == @on_input && @input.eof?
     end
   end
 end

data/lib/kronk/queue_runner.rb

@@ -0,0 +1,238 @@
+class Kronk
+
+  class QueueRunner
+
+    ##
+    # Define whether to use the EventMachine or the threaded behavior.
+
+    def self.async= value
+      @async = !!value
+    end
+
+
+    ##
+    # Returns true if EventMachine is enabled
+
+    def self.async
+      @async
+    end
+
+    self.async = false
+
+
+    attr_accessor :number, :concurrency, :queue, :count,
+                  :mutex, :threads, :reader_thread
+
+    ##
+    # Create a new QueueRunner for batch multi-threaded processing.
+    # Supported options are:
+    # :concurrency:: Fixnum - Maximum number of concurrent items to process
+    # :number:: Fixnum - Total number of items to process
+
+    def initialize opts={}
+      @number      = opts[:number]
+      @concurrency = opts[:concurrency]
+      @concurrency = 1 if !@concurrency || @concurrency <= 0
+
+      @count   = 0
+      @queue   = []
+      @threads = []
+
+      @reader_thread = nil
+
+      @triggers = {}
+
+      @mutex = Mutex.new
+    end
+
+
+    ##
+    # Returns true if processing queue should be stopped, otherwise false.
+
+    def finished?
+      return true if @number && @count >= @number
+
+      @queue.empty? && @count > 0 &&
+        (!@reader_thread || !@reader_thread.alive?)
+    end
+
+
+    ##
+    # Immediately end all runner processing and threads.
+
+    def kill
+      stop_input!
+      EM.stop if defined?(EM) && EM.reactor_running?
+      @threads.each{|t| t.kill}
+      @threads.clear
+    end
+
+
+    ##
+    # Specify a block to run for a given trigger name.
+    # Supported triggers are:
+    # :complete:: Called after queue and input have been fully processed.
+    # :input:: Called every time the queue needs populating.
+    # :interrupt:: Called when SIGINT is captured.
+    # :start:: Called before queue starts being processed.
+
+    def on trigger_name, &block
+      @triggers[trigger_name] = block
+    end
+
+
+    ##
+    # Start processing the queue and reading from IO if available.
+    #
+    # Yields queue item until queue and io (if available) are empty and the
+    # totaly number of requests to run is met (if number is set).
+
+    def process_queue
+      start_input!
+      @count = 0
+
+      until finished?
+        @threads.delete_if{|t| !t.alive? }
+
+        if @threads.length >= @concurrency || @queue.empty?
+          Thread.pass
+          next
+        end
+
+        @threads << Thread.new(@queue.shift) do |q_item|
+          yield q_item if block_given?
+        end
+
+        @count += 1
+      end
+
+      @threads.each{|t| t.join}
+      @threads.clear
+
+      stop_input!
+    end
+
+
+    ##
+    # Start processing the queue and reading from IO if available.
+    #
+    # Yields queue item until queue and io (if available) are empty and the
+    # totaly number of requests to run is met (if number is set).
+    #
+    # Uses EventMachine to run asynchronously.
+    #
+    # Note: If the block given doesn't use EM, it will be blocking.
+
+    def process_queue_async &block
+      # TODO: Make input use EM from QueueRunner and Player IO.
+      require 'kronk/async' unless defined?(EM::HttpRequest)
+      Cmd.verbose "Running async"
+
+      start_input!
+
+      @count = 0
+
+      EM.run do
+        EM.add_periodic_timer do
+          if finished?
+            next if EM.connection_count > 0
+            kill
+            next
+          end
+
+          if @queue.empty? || EM.connection_count >= @concurrency
+            Thread.pass
+            next
+          end
+
+          yield @queue.shift
+          @count += 1
+        end
+      end
+    end
+
+
+    ##
+    # Runs the queue and reads from input until it's exhausted or
+    # @number is reached. Yields a queue item and a mutex when to passed
+    # block:
+    #
+    #   runner = QueueRunner.new :concurrency => 10
+    #   runner.queue.concat %w{item1 item2 item3}
+    #
+    #   runner.run do |q_item, mutex|
+    #     # This block is run in its own thread.
+    #     mutex.synchronize{ do_something_with q_item }
+    #   end
+    #
+    # Calls the :start trigger before execution begins, calls :complete
+    # when the execution has ended or is interrupted, also calls :interrupt
+    # when execution is interrupted.
+
+    def run
+      trap 'INT' do
+        kill
+        (trigger(:interrupt) || exit(1))
+      end
+
+      trigger :start
+
+      method = self.class.async ? :process_queue_async : :process_queue
+
+      send method do |q_item|
+        yield q_item, @mutex if block_given?
+      end
+
+      trigger :complete
+    end
+
+
+    ##
+    # Attempt to fill the queue by reading from the IO instance.
+    # Starts a new thread and returns the thread instance.
+
+    def start_input!
+      return unless @triggers[:input]
+
+      max_queue_size = @concurrency * 2
+
+      @reader_thread = Thread.new do
+        begin
+          loop do
+            if @queue.length >= max_queue_size
+              Thread.pass
+              next
+            end
+
+            while @queue.length < max_queue_size
+              @queue << trigger(:input)
+            end
+            Thread.pass
+          end
+
+        rescue => e
+          Thread.main.raise e
+        end
+      end
+    end
+
+
+    ##
+    # Permanently stop input reading by killing the reader thread for a given
+    # QueueRunner#run or QueueRunner#process_queue session.
+
+    def stop_input!
+      Thread.pass
+      @reader_thread && @reader_thread.kill
+    end
+
+
+    ##
+    # Run a previously defined callback. See QueueRunner#on.
+
+    def trigger name
+      t = @triggers[name]
+      t && t.call
+    end
+  end
+end

data/lib/kronk/request.rb

@@ -1,15 +1,17 @@
 class Kronk
 
   ##
-  # Performs HTTP requests or retrieves HTTP responses.
+  # Performs HTTP requests and returns a Kronk::Response instance.
 
   class Request
 
     # Raised by Request.parse when parsing invalid http request string.
     class ParseError < Kronk::Exception; end
 
-    # Matches the first line of an http request string.
-    REQUEST_LINE_MATCHER = %r{([A-Za-z]+)?(^|[\s'"])(/[^\s'";]+)[\s"']*}
+    # Matches the first line of an http request string or a fully
+    # qualified URL.
+    REQUEST_LINE_MATCHER =
+ %r{(?:^|[\s'"])(?:([a-z]+)\s)?(?:(https?://[^/]+)(/[^\s'";]*)?|(/[^\s'";]*))}i
 
     ##
     # Creates a query string from data.
@@ -45,7 +47,7 @@ class Kronk
     # path and options.
 
     def self.build_uri uri, options={}
-      uri  ||= Kronk.config[:default_host]
+      uri  ||= options[:host] || Kronk.config[:default_host]
       suffix = options[:uri_suffix]
 
       uri = "http://#{uri}"   unless uri.to_s =~ %r{^(\w+://|/)}
@@ -64,6 +66,7 @@ class Kronk
 
     ##
     # Parses a raw HTTP request-like string into a Kronk::Request instance.
+    # Options passed are used as default values for Request#new.
 
     def self.parse str, opts={}
       opts = parse_to_hash str, opts
@@ -75,7 +78,8 @@ class Kronk
 
     ##
     # Parses a raw HTTP request-like string into a Kronk::Request options hash.
-    # Also parses most single access log entries.
+    # Also parses most single access log entries. Options passed are used
+    # as default values for Request#new.
 
     def self.parse_to_hash str, opts={}
       lines = str.split("\n")
@@ -86,7 +90,9 @@ class Kronk
       opts[:headers] ||= {}
 
       lines.shift.strip =~ REQUEST_LINE_MATCHER
-      opts[:http_method], opts[:uri_suffix] = $1, $3
+      opts.merge! :http_method => $1,
+                  :host        => $2,
+                  :uri_suffix  => ($3 || $4)
 
       lines.each_with_index do |line, i|
         case line
@@ -105,6 +111,7 @@ class Kronk
 
       opts[:data] = lines[body_start..-1].join("\n") if body_start
 
+      opts.delete(:host)        if !opts[:host]
       opts.delete(:uri_suffix)  if !opts[:uri_suffix]
       opts.delete(:headers)     if opts[:headers].empty?
       opts.delete(:http_method) if !opts[:http_method]
@@ -178,7 +185,7 @@ class Kronk
     end
 
 
-    attr_accessor :body, :headers, :response, :timeout
+    attr_accessor :body, :headers, :proxy, :response, :timeout
 
     attr_reader :http_method, :uri, :use_cookies
 
@@ -187,18 +194,16 @@ class Kronk
     # Supports the following options:
     # :data:: Hash/String - the data to pass to the http request
     # :query:: Hash/String - the data to append to the http request path
-    # :follow_redirects:: Integer/Bool - number of times to follow redirects
     # :user_agent:: String - user agent string or alias; defaults to 'kronk'
     # :auth:: Hash - must contain :username and :password; defaults to nil
     # :headers:: Hash - extra headers to pass to the request
     # :http_method:: Symbol - the http method to use; defaults to :get
-    # :proxy:: Hash/String - http proxy to use; defaults to nil
+    # :proxy:: Hash/String - http proxy to use; defaults to {}
     #
     # Note: if no http method is specified and data is given, will default
     # to using a post request.
 
     def initialize uri, options={}
-      @HTTP = Net::HTTP
       @auth = options[:auth]
 
       @body = nil
@@ -213,18 +218,15 @@ class Kronk
 
       @uri = self.class.build_uri uri, options
 
+      @proxy = options[:proxy] || {}
+      @proxy = {:host => @proxy} unless Hash === @proxy
+
       self.user_agent ||= options[:user_agent]
 
       self.http_method = options[:http_method] || (@body ? "POST" : "GET")
 
       self.use_cookies = options.has_key?(:no_cookies) ?
                           !options[:no_cookies] : Kronk.config[:use_cookies]
-
-      if Hash === options[:proxy]
-        self.use_proxy options[:proxy][:address], options[:proxy]
-      else
-        self.use_proxy options[:proxy]
-      end
     end
 
 
@@ -278,8 +280,8 @@ class Kronk
     # The proxy_opts arg can be a uri String or a Hash with the :address key
     # and optional :username and :password keys.
 
-    def use_proxy addr, opts={}
-      return @HTTP = Net::HTTP unless addr
+    def http_proxy addr, opts={}
+      return Net::HTTP unless addr
 
       host, port = addr.split ":"
       port ||= opts[:port] || 8080
@@ -289,7 +291,7 @@ class Kronk
 
       Kronk::Cmd.verbose "Using proxy #{addr}\n" if host
 
-      @HTTP = Net::HTTP::Proxy host, port, user, pass
+      Net::HTTP::Proxy host, port, user, pass
     end
 
 
@@ -355,7 +357,9 @@ class Kronk
     # Retrieve this requests' response.
 
     def retrieve
-      @_req = @HTTP.new @uri.host, @uri.port
+      http_class = http_proxy @proxy[:host], @proxy
+
+      @_req = http_class.new @uri.host, @uri.port
 
       @_req.read_timeout = @timeout if @timeout
       @_req.use_ssl      = true     if @uri.scheme =~ /^https$/
@@ -424,3 +428,4 @@ class Kronk
     end
   end
 end
+