buildkite-test_collector 1.0.0

data/lib/buildkite/test_collector/session.rb

@@ -0,0 +1,331 @@
+# frozen_string_literal: true
+
+require_relative "socket_connection"
+
+module Buildkite::TestCollector
+  class Session
+    # Picked 75 as the magic timeout number as it's longer than the TCP timeout of 60s 🤷‍♀️
+    CONFIRMATION_TIMEOUT = ENV.fetch("BUILDKITE_ANALYTICS_CONFIRMATION_TIMEOUT") { 75 }.to_i
+    MAX_RECONNECTION_ATTEMPTS = ENV.fetch("BUILDKITE_ANALYTICS_RECONNECTION_ATTEMPTS") { 3 }.to_i
+    WAIT_BETWEEN_RECONNECTIONS = ENV.fetch("BUILDKITE_ANALYTICS_RECONNECTION_WAIT") { 5 }.to_i
+
+    class RejectedSubscription < StandardError; end
+    class InitialConnectionFailure < StandardError; end
+
+    DISCONNECTED_EXCEPTIONS = [
+      SocketConnection::HandshakeError,
+      RejectedSubscription,
+      TimeoutError,
+      InitialConnectionFailure,
+      SocketConnection::SocketError
+    ]
+
+    def initialize(url, authorization_header, channel)
+      @establish_subscription_queue = Queue.new
+      @channel = channel
+
+      @unconfirmed_idents = {}
+      @idents_mutex = Mutex.new
+      @send_queue = Queue.new
+      @empty = ConditionVariable.new
+      @closing = false
+      @eot_queued = false
+      @eot_queued_mutex = Mutex.new
+      @reconnection_mutex = Mutex.new
+
+      @url = url
+      @authorization_header = authorization_header
+
+      reconnection_count = 0
+
+      begin
+        reconnection_count += 1
+        connect
+      rescue TimeoutError, InitialConnectionFailure => e
+        Buildkite::TestCollector.logger.warn("rspec-buildkite-analytics could not establish an initial connection with Buildkite due to #{e}. Attempting retry #{reconnection_count} of #{MAX_RECONNECTION_ATTEMPTS}...")
+        if reconnection_count > MAX_RECONNECTION_ATTEMPTS
+          Buildkite::TestCollector.logger.error "rspec-buildkite-analytics could not establish an initial connection with Buildkite due to #{e.message} after #{MAX_RECONNECTION_ATTEMPTS} attempts. You may be missing some data for this test suite, please contact support if this issue persists."
+        else
+          sleep(WAIT_BETWEEN_RECONNECTIONS)
+          Buildkite::TestCollector.logger.warn("retrying reconnection")
+          retry
+        end
+      end
+      init_write_thread
+    end
+
+    def disconnected(connection)
+      @reconnection_mutex.synchronize do
+        # When the first thread detects a disconnection, it calls the disconnect method
+        # with the current connection. This thread grabs the reconnection mutex and does the
+        # reconnection, which then updates the value of @connection.
+        #
+        # At some point in that process, the second thread would have detected the
+        # disconnection too, and it also calls it with the current connection. However, the
+        # second thread can't run the reconnection code because of the mutex. By the
+        # time the mutex is released, the value of @connection has been refreshed, and so
+        # the second thread returns early and does not reattempt the reconnection.
+        return unless connection == @connection
+        Buildkite::TestCollector.logger.debug("starting reconnection")
+
+        reconnection_count = 0
+
+        begin
+          reconnection_count += 1
+          connect
+          init_write_thread
+        rescue *DISCONNECTED_EXCEPTIONS => e
+          Buildkite::TestCollector.logger.warn("failed reconnection attempt #{reconnection_count} due to #{e}")
+          if reconnection_count > MAX_RECONNECTION_ATTEMPTS
+            Buildkite::TestCollector.logger.error "rspec-buildkite-analytics experienced a disconnection and could not reconnect to Buildkite due to #{e.message}. Please contact support."
+            raise e
+          else
+            sleep(WAIT_BETWEEN_RECONNECTIONS)
+            Buildkite::TestCollector.logger.warn("retrying reconnection")
+            retry
+          end
+        end
+      end
+      retransmit
+    end
+
+    def close(examples_count)
+      @closing = true
+      @examples_count = examples_count
+      Buildkite::TestCollector.logger.debug("closing socket connection")
+
+      # Because the server only sends us confirmations after every 10mb of
+      # data it uploads to S3, we'll never get confirmation of the
+      # identifiers of the last upload part unless we send an explicit finish,
+      # to which the server will respond with the last bits of data
+      send_eot
+
+      # After EOT, we wait for 75 seconds for the send queue to be drained and for the
+      # server to confirm the last idents. If everything has already been confirmed we can
+      # proceed without waiting.
+      @idents_mutex.synchronize do
+        if @unconfirmed_idents.any?
+          Buildkite::TestCollector.logger.debug "Waiting for Buildkite Test Analytics to send results..."
+          Buildkite::TestCollector.logger.debug("waiting for last confirm")
+
+          @empty.wait(@idents_mutex, CONFIRMATION_TIMEOUT)
+        end
+      end
+
+      # Then we always disconnect cos we can't wait forever? 🤷‍♀️
+      @connection.close
+      # We kill the write thread cos it's got a while loop in it, so it won't finish otherwise
+      @write_thread&.kill
+
+      Buildkite::TestCollector.logger.info "Buildkite Test Analytics completed"
+      Buildkite::TestCollector.logger.debug("socket connection closed")
+    end
+
+    def handle(_connection, data)
+      data = JSON.parse(data)
+      case data["type"]
+      when "ping"
+        # In absence of other message, the server sends us a ping every 3 seconds
+        # We are currently not doing anything with these
+        Buildkite::TestCollector.logger.debug("received ping")
+      when "welcome", "confirm_subscription"
+        # Push these two messages onto the queue, so that we block on waiting for the
+        # initializing phase to complete
+        @establish_subscription_queue.push(data)
+      Buildkite::TestCollector.logger.debug("received #{data['type']}")
+      when "reject_subscription"
+        Buildkite::TestCollector.logger.debug("received rejected_subscription")
+        raise RejectedSubscription
+      else
+        process_message(data)
+      end
+    end
+
+    def write_result(result)
+      queue_and_track_result(result.id, result.as_hash)
+
+      Buildkite::TestCollector.logger.debug("added #{result.id} to send queue")
+    end
+
+    def unconfirmed_idents_count
+      @idents_mutex.synchronize do
+        @unconfirmed_idents.count
+      end
+    end
+
+    private
+
+    def connect
+      Buildkite::TestCollector.logger.debug("starting socket connection process")
+
+      @connection = SocketConnection.new(self, @url, {
+        "Authorization" => @authorization_header,
+      })
+
+      wait_for_welcome
+
+      @connection.transmit({
+        "command" => "subscribe",
+        "identifier" => @channel
+      })
+
+      wait_for_confirm
+
+      Buildkite::TestCollector.logger.info "Connected to Buildkite Test Analytics!"
+      Buildkite::TestCollector.logger.debug("connected")
+    end
+
+    def init_write_thread
+      # As this method can be called multiple times in the
+      # reconnection process, kill prev write threads (if any) before
+      # setting up the new one
+      @write_thread&.kill
+
+      @write_thread = Thread.new do
+        Buildkite::TestCollector.logger.debug("hello from write thread")
+        # Pretty sure this eternal loop is fine cos the call to queue.pop is blocking
+        loop do
+          data = @send_queue.pop
+          message_type = data["action"]
+
+          if message_type == "end_of_transmission"
+            # Because of the unpredictable sequencing between the test suite finishing
+            # (EOT gets queued) and disconnections happening (retransmit results gets
+            # queued), we don't want to send an EOT before any retransmits are sent.
+            if @send_queue.length > 0
+              @send_queue << data
+              Buildkite::TestCollector.logger.debug("putting eot at back of queue")
+              next
+            end
+            @eot_queued_mutex.synchronize do
+              @eot_queued = false
+            end
+          end
+
+          @connection.transmit({
+            "identifier" => @channel,
+            "command" => "message",
+            "data" => data.to_json
+          })
+
+          if Buildkite::TestCollector.debug_enabled
+            ids = if message_type == "record_results"
+              data["results"].map { |result| result["id"] }
+            end
+            Buildkite::TestCollector.logger.debug("transmitted #{message_type} #{ids}")
+          end
+        end
+      end
+    end
+
+    def pop_with_timeout(message_type)
+      Timeout.timeout(30, Buildkite::TestCollector::TimeoutError, "Timeout: Waited 30 seconds for #{message_type}") do
+        @establish_subscription_queue.pop
+      end
+    end
+
+    def wait_for_welcome
+      welcome = pop_with_timeout("welcome")
+
+      if welcome && welcome != { "type" => "welcome" }
+        raise InitialConnectionFailure.new("Wrong message received, expected a welcome, but received: #{welcome.inspect}")
+      end
+    end
+
+    def wait_for_confirm
+      confirm = pop_with_timeout("confirm")
+
+      if confirm && confirm != { "type" => "confirm_subscription", "identifier" => @channel }
+        raise InitialConnectionFailure.new("Wrong message received, expected a confirm, but received: #{confirm.inspect}")
+      end
+    end
+
+    def queue_and_track_result(ident, result_as_hash)
+      @idents_mutex.synchronize do
+        @unconfirmed_idents[ident] = result_as_hash
+
+        @send_queue << {
+          "action" => "record_results",
+          "results" => [result_as_hash]
+        }
+      end
+    end
+
+    def confirm_idents(idents)
+      retransmit_required = @closing
+
+      @idents_mutex.synchronize do
+        # Remove received idents from unconfirmed_idents
+        idents.each { |key| @unconfirmed_idents.delete(key) }
+
+        Buildkite::TestCollector.logger.debug("received confirm for indentifiers: #{idents}")
+
+        # This @empty ConditionVariable broadcasts every time that @unconfirmed_idents is
+        # empty, which will happen about every 10mb of data as that's when the server
+        # sends back confirmations.
+        #
+        # However, there aren't any threads waiting on this signal until after we
+        # send the EOT message, so the prior broadcasts shouldn't do anything.
+        if @unconfirmed_idents.empty?
+          @empty.broadcast
+
+          retransmit_required = false
+
+          Buildkite::TestCollector.logger.debug("all identifiers have been confirmed")
+        else
+          Buildkite::TestCollector.logger.debug("still waiting on confirm for identifiers: #{@unconfirmed_idents.keys}")
+        end
+      end
+
+      # If we're closing, any unconfirmed results need to be retransmitted.
+      retransmit if retransmit_required
+    end
+
+    def send_eot
+      @eot_queued_mutex.synchronize do
+        return if @eot_queued
+
+        @send_queue << {
+          "action" => "end_of_transmission",
+          "examples_count" => @examples_count.to_json
+        }
+        @eot_queued = true
+
+        Buildkite::TestCollector.logger.debug("added EOT to send queue")
+      end
+    end
+
+    def process_message(data)
+      # Check we're getting the data we expect
+      return unless data["identifier"] == @channel
+
+      case
+      when data["message"].key?("confirm")
+        confirm_idents(data["message"]["confirm"])
+      else
+        # unhandled message
+        Buildkite::TestCollector.logger.debug("received unhandled message #{data["message"]}")
+      end
+    end
+
+    def retransmit
+      @idents_mutex.synchronize do
+        results = @unconfirmed_idents.values
+
+        # queue the contents of the buffer, unless it's empty
+        if results.any?
+          @send_queue << {
+            "action" => "record_results",
+            "results" => results
+          }
+
+          Buildkite::TestCollector.logger.debug("queueing up retransmitted results #{@unconfirmed_idents.keys}")
+        end
+      end
+
+      # if we were disconnected in the closing phase, then resend the EOT
+      # message so the server can persist the last upload part
+      send_eot if @closing
+    end
+  end
+end

data/lib/buildkite/test_collector/socket_connection.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require "socket"
+require "openssl"
+require "json"
+
+module Buildkite::TestCollector
+  class SocketConnection
+    class HandshakeError < StandardError; end
+    class SocketError < StandardError; end
+
+    def initialize(session, url, headers)
+      uri = URI.parse(url)
+      @session = session
+      protocol = "http"
+
+      begin
+        socket = TCPSocket.new(uri.host, uri.port || (uri.scheme == "wss" ? 443 : 80))
+
+        if uri.scheme == "wss"
+          ctx = OpenSSL::SSL::SSLContext.new
+          protocol = "https"
+
+          ctx.min_version = :TLS1_2
+          ctx.verify_mode = OpenSSL::SSL::VERIFY_PEER
+          ctx.cert_store = OpenSSL::X509::Store.new.tap(&:set_default_paths)
+
+          socket = OpenSSL::SSL::SSLSocket.new(socket, ctx)
+          socket.connect
+        end
+      rescue
+        # We are rescuing all here, as there are a range of Errno errors that could be
+        # raised when we fail to establish a TCP connection
+        raise SocketError
+      end
+
+      @socket = socket
+
+      headers = { "Origin" => "#{protocol}://#{uri.host}" }.merge(headers)
+      handshake = WebSocket::Handshake::Client.new(url: url, headers: headers)
+
+      @socket.write handshake.to_s
+
+      until handshake.finished?
+        if byte = @socket.getc
+          handshake << byte
+        end
+      end
+
+      # The errors below are raised when we establish the TCP connection, but get back
+      # an error, i.e. in dev we can still connect to puma-dev while nginx isn't
+      # running, or in prod we can hit a load balancer while app is down
+      unless handshake.valid?
+        case handshake.error
+        when Exception, String
+          raise HandshakeError.new(handshake.error)
+        when nil
+          raise HandshakeError.new("Invalid handshake")
+        else
+          raise HandshakeError.new(handshake.error.inspect)
+        end
+      end
+
+      @version = handshake.version
+
+      # Setting up a new thread that listens on the socket, and processes incoming
+      # comms from the server
+      @thread = Thread.new do
+        Buildkite::TestCollector.logger.debug("listening in on socket")
+        frame = WebSocket::Frame::Incoming::Client.new
+
+        while @socket
+          frame << @socket.readpartial(4096)
+
+          while data = frame.next
+            @session.handle(self, data.data)
+          end
+        end
+      # These get re-raise from session, we should fail gracefully
+      rescue *Buildkite::TestCollector::Session::DISCONNECTED_EXCEPTIONS => e
+        Buildkite::TestCollector.logger.error("We could not establish a connection with Buildkite Test Analytics. The error was: #{e.message}. If this is a problem, please contact support.")
+      rescue EOFError => e
+        Buildkite::TestCollector.logger.warn("#{e}")
+        if @socket
+          Buildkite::TestCollector.logger.warn("attempting disconnected flow")
+          @session.disconnected(self)
+          disconnect
+        end
+      rescue Errno::ECONNRESET => e
+        Buildkite::TestCollector.logger.error("#{e}")
+        if @socket
+          Buildkite::TestCollector.logger.error("attempting disconnected flow")
+          @session.disconnected(self)
+          disconnect
+        end
+      rescue IOError
+        # This is fine to ignore
+        Buildkite::TestCollector.logger.error("IOError")
+      rescue IndexError
+        # I don't like that we're doing this but I think it's the best of the options
+        #
+        # This relates to this issue https://github.com/ruby/openssl/issues/452
+        # A fix for it has been released but the repercussions of overriding
+        # the OpenSSL version in the stdlib seem worse than catching this error here.
+        Buildkite::TestCollector.logger.error("IndexError")
+        if @socket
+          Buildkite::TestCollector.logger.error("attempting disconnected flow")
+          @session.disconnected(self)
+          disconnect
+        end
+      end
+    end
+
+    def transmit(data, type: :text)
+      # this line prevents us from calling disconnect twice
+      return if @socket.nil?
+
+      raw_data = data.to_json
+      frame = WebSocket::Frame::Outgoing::Client.new(data: raw_data, type: :text, version: @version)
+      @socket.write(frame.to_s)
+    rescue Errno::EPIPE, Errno::ECONNRESET, OpenSSL::SSL::SSLError => e
+      return unless @socket
+      return if type == :close
+      Buildkite::TestCollector.logger.error("got #{e}, attempting disconnected flow")
+      @session.disconnected(self)
+      disconnect
+    rescue IndexError
+      # I don't like that we're doing this but I think it's the best of the options
+      #
+      # This relates to this issue https://github.com/ruby/openssl/issues/452
+      # A fix for it has been released but the repercussions of overriding
+      # the OpenSSL version in the stdlib seem worse than catching this error here.
+      Buildkite::TestCollector.logger.error("IndexError")
+      if @socket
+        Buildkite::TestCollector.logger.error("attempting disconnected flow")
+        @session.disconnected(self)
+        disconnect
+      end
+    end
+
+    def close
+      Buildkite::TestCollector.logger.debug("socket close")
+      transmit(nil, type: :close)
+      disconnect
+    end
+
+    private
+
+    def disconnect
+      Buildkite::TestCollector.logger.debug("socket disconnect")
+      socket = @socket
+      @socket = nil
+      socket&.close
+      @thread&.join unless @thread == Thread.current
+    end
+  end
+end

data/lib/buildkite/test_collector/tracer.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/hash/indifferent_access"
+
+module Buildkite::TestCollector
+  class Tracer
+    class Span
+      attr_accessor :section, :start_at, :end_at, :detail, :children
+
+      def initialize(section, start_at, end_at, detail)
+        @section = section
+        @start_at = start_at
+        @end_at = end_at
+        @detail = detail
+        @children = []
+      end
+
+      def as_hash
+        {
+          section: section,
+          start_at: start_at,
+          end_at: end_at,
+          duration: end_at - start_at,
+          detail: detail,
+          children: children.map(&:as_hash),
+        }.with_indifferent_access
+      end
+    end
+
+    def initialize
+      @top = Span.new(:top, Concurrent.monotonic_time, nil, {})
+      @stack = [@top]
+    end
+
+    def enter(section, **detail)
+      new_entry = Span.new(section, Concurrent.monotonic_time, nil, detail)
+      current_span.children << new_entry
+      @stack << new_entry
+    end
+
+    def leave
+      current_span.end_at = Concurrent.monotonic_time
+      @stack.pop
+    end
+
+    def backfill(section, duration, **detail)
+      new_entry = Span.new(section, Concurrent.monotonic_time - duration, Concurrent.monotonic_time, detail)
+      current_span.children << new_entry
+    end
+
+    def current_span
+      @stack.last
+    end
+
+    def finalize
+      raise "Stack not empty" unless @stack.size == 1
+      @top.end_at = Concurrent.monotonic_time
+      self
+    end
+
+    def history
+      @top.as_hash
+    end
+  end
+end

data/lib/buildkite/test_collector/uploader.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "openssl"
+require "websocket"
+
+require_relative "tracer"
+require_relative "network"
+require_relative "object"
+require_relative "session"
+require_relative "ci"
+require_relative "http_client"
+
+require "active_support"
+require "active_support/notifications"
+
+require "securerandom"
+
+module Buildkite::TestCollector
+  class Uploader
+    def self.traces
+      @traces ||= {}
+    end
+
+    REQUEST_EXCEPTIONS = [
+      URI::InvalidURIError,
+      Net::HTTPBadResponse,
+      Net::HTTPHeaderSyntaxError,
+      Net::ReadTimeout,
+      Net::OpenTimeout,
+      OpenSSL::SSL::SSLError,
+      OpenSSL::SSL::SSLErrorWaitReadable,
+      EOFError
+    ]
+
+    def self.configure
+      Buildkite::TestCollector.logger.debug("hello from main thread")
+
+      if Buildkite::TestCollector.api_token
+        http = Buildkite::TestCollector::HTTPClient.new(Buildkite::TestCollector.url)
+
+        response = begin
+          http.post
+        rescue *Buildkite::TestCollector::Uploader::REQUEST_EXCEPTIONS => e
+          Buildkite::TestCollector.logger.error "Buildkite Test Analytics: Error communicating with the server: #{e.message}"
+        end
+
+        return unless response
+
+        case response.code
+        when "401"
+          Buildkite::TestCollector.logger.info "Buildkite Test Analytics: Invalid Suite API key. Please double check your Suite API key."
+        when "200"
+          json = JSON.parse(response.body)
+
+          if (socket_url = json["cable"]) && (channel = json["channel"])
+            Buildkite::TestCollector.session = Buildkite::TestCollector::Session.new(socket_url, http.authorization_header, channel)
+          end
+        else
+          request_id = response.to_hash["x-request-id"]
+          Buildkite::TestCollector.logger.info "rspec-buildkite-analytics could not establish an initial connection with Buildkite. You may be missing some data for this test suite, please contact support."
+        end
+      else
+        if !!ENV["BUILDKITE_BUILD_ID"]
+          Buildkite::TestCollector.logger.info "Buildkite Test Analytics: No Suite API key provided. You can get the API key from your Suite settings page."
+        end
+      end
+    end
+
+    def self.tracer
+      Thread.current[:_buildkite_tracer]
+    end
+  end
+end

data/lib/buildkite/test_collector/version.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module Buildkite
+  module TestCollector
+    VERSION = "1.0.0"
+    NAME = "buildkite-test_collector"
+  end
+end