HTTPS load testing

load.rb

require "socket"
require "openssl"
require "thread"
require "thwait"
require "benchmark"
require "optparse"
require "net/http"
require "csv"
require "uri"

class Request
  METRICS = [:connection_time, :handshake_time, :send_time, :receive_time]

  attr_accessor :started_at, :ended_at, :id, :thread_id, :response, :error

  def status_code
    response.code.to_i if response
  end

  def headers
    response.to_hash if response
  end

  def finished?
    !ended_at.nil?
  end

  def success?
    finished? && response && status_code >= 200 && status_code < 400
  end

  def failed?
    finished? && (error || !response || status_code >= 400)
  end

  METRICS.each do |metric|
    define_method("#{metric}") do
      instance_variable_get("@#{metric}").to_f * 1000
    end

    define_method("measure_#{metric}") do |&block|
      instance_variable_set "@#{metric}", Benchmark.realtime(&block)
    end
  end

  def response_time
    METRICS.map { |metric| send(metric) }.reduce(:+)
  end
end

class Report < Array
  (Request::METRICS + [:response_time]).each do |metric|
    define_method("mean_#{metric}") do
      return 0 if empty?
      map { |request| request.send(metric) }.select { |value| value > 0 }.reduce(:+) / size
    end

    define_method("max_#{metric}") do
      return 0 if empty?
      map { |request| request.send(metric) }.select { |value| value > 0 }.max
    end

    define_method("min_#{metric}") do
      return 0 if empty?
      map { |request| request.send(metric) }.select { |value| value > 0 }.min
    end
  end

  def error_rate
    return 0 if empty?
    (count(&:failed?) / size.to_f) * 100
  end

  def elapsed_time
    return 0 if finished_requests.empty?
    finished_requests.map(&:ended_at).max - finished_requests.map(&:started_at).min
  end

  def finished_requests
    select(&:finished?)
  end

  def to_s
    [
      ["Requests", size],
      ["Error rate", format_decimal(error_rate) + "%"],
      [],

      ["Avg response time (ms)", format_decimal(mean_response_time)],
      ["Min response time (ms)", format_decimal(min_response_time)],
      ["Max response time (ms)", format_decimal(max_response_time)],
      [],

      ["Avg connection time (ms)", format_decimal(mean_connection_time)],
      ["Min connection time (ms)", format_decimal(min_connection_time)],
      ["Max connection time (ms)", format_decimal(max_connection_time)],
      [],

      ["Avg handshake time (ms)", format_decimal(mean_handshake_time)],
      ["Min handshake time (ms)", format_decimal(min_handshake_time)],
      ["Max handshake time (ms)", format_decimal(max_handshake_time)],
      [],

      ["Avg send time (ms)", format_decimal(mean_send_time)],
      ["Min send time (ms)", format_decimal(min_send_time)],
      ["Max send time (ms)", format_decimal(max_send_time)],
      [],

      ["Avg receive time (ms)", format_decimal(mean_receive_time)],
      ["Min receive time (ms)", format_decimal(min_receive_time)],
      ["Max receive time (ms)", format_decimal(max_receive_time)],
    ].map { |line| line.join(": ") }.join("\n")
  end

  def to_csv
    def format(number)
      format_decimal(number, separator: ",")
    end

    CSV.generate do |csv|
      each do |request|
        csv << [
          format(request.response_time),
          format(request.connection_time),
          format(request.handshake_time),
          format(request.send_time),
          format(request.receive_time),
          request.id,
          request.thread_id,
          request.status_code,
          request.error
        ]
      end
    end
  end

  private

  def format_decimal(number, separator: ".")
    number.round(2).to_s.sub(".", separator)
  end
end

class Runner
  def initialize(concurrency: 10, requests: 100, keep_alive: false)
    @concurrency = concurrency
    @requests = requests
    @keep_alive = keep_alive
  end

  def run(http_method, url, headers, &callback)
    url = URI.parse(url)
    report = Report.new
    results = Queue.new
    expected_requests = @requests * @concurrency
    requests_done = 0

    ready_threads = 0
    should_start = false
    mutex = Mutex.new

    http_request = "#{http_method} #{url.path} HTTP/1.1\r\n"
    http_request << "Host: #{url.host}\r\n"
    http_request << headers.join("\r\n")
    http_request << "\n\r\n\r\n\r"

    Thread.abort_on_exception = true

    threads = @concurrency.times.map do |thread_id|
      Thread.new(@concurrency, @requests, @keep_alive) do |concurrency, requests, keep_alive|
        mutex.synchronize do
          ready_threads += 1
          should_start = ready_threads == concurrency
        end
        sleep 0.1 until should_start

        socket = ssl = nil

        requests.times do |request_number|
          request = Request.new
          request.thread_id = thread_id + 1
          request.id = request_number * request.thread_id

          begin
            socket = ssl = nil unless keep_alive

            request.started_at = Time.now
            request.measure_connection_time do
              socket = TCPSocket.new(url.host, url.port) if socket.nil?
            end

            if ssl.nil?
              context = OpenSSL::SSL::SSLContext.new
              context.set_params(verify_mode: OpenSSL::SSL::VERIFY_NONE)

              ssl = OpenSSL::SSL::SSLSocket.new(socket, context)
              ssl.hostname = url.host
              ssl.sync_close = true
              request.measure_handshake_time { ssl.connect }
            end

            request.measure_send_time { ssl.write(http_request) }

            request.measure_receive_time do
              request.response = Net::HTTPResponse.read_new(Net::BufferedIO.new(ssl))
            end

            ssl.close unless keep_alive
          rescue StandardError => error
            request.error = error
          ensure
            request.ended_at = Time.now
            results << request
            mutex.synchronize { requests_done += 1 }
          end
        end

        ssl.close if keep_alive
      end
    end

    threads << Thread.new do
      sleep 0.1 until should_start

      loop do
        if !results.empty?
          report << results.pop until results.empty?
          callback.call(report) if callback
        end

        break if expected_requests == requests_done

        sleep 0.1
      end
    end

    ThreadsWait.all_waits(*threads)

    report
  end
end

options = {}
output_csv = false
http_method = "HEAD"
headers = []

OptionParser.new do |opts|
  opts.banner = "Usage: load.rb [options] url"

  opts.on("-X", "--method M", String, "HTTP method") do |v|
    http_method = v
  end

  opts.on("-H", "--header H", String, "HTTP header") do |v|
    headers << v
  end

  opts.on("-n", "--number N", Integer, "Requests") do |v|
    options[:requests] = v
  end

  opts.on("-c", "--concurrency N", Integer, "Concurrency") do |v|
    options[:concurrency] = v
  end

  opts.on("-k", "--keep-alive", "Doesn't close connections") do |v|
    options[:keep_alive] = v
  end

  opts.on("--csv", "CSV output") do |v|
    output_csv = v
  end
end.parse!

url = ARGV.last or abort("You must inform a url")

runner = Runner.new(**options)

report = runner.run(http_method, url, headers) do |report|
  output = report.to_s
  output.lines.count.times { STDERR.print "\r\e[A\e[K" }
  STDERR.puts output
end

puts report.to_csv if output_csv