buildkite-test_collector 1.4.2 → 2.0.0.pre

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c89588cb63374f9a2bd84b29cd575133aa84c2a684e10e0e7d78dd63a54ad34e
-  data.tar.gz: a05196237e4d03ae61f9953c94de3d54253a40c89a6f0557c3fa8e42b800e163
+  metadata.gz: b7fc00254b8a2002e8680ca265b00f5b2700dace5d6948aa56f779102f8ce0c6
+  data.tar.gz: 6c05b872c740262ca0c66f4d6c09b5b253a366466241513b74706c8d0a56328a
 SHA512:
-  metadata.gz: ab0e5a4fc6a21adc7f8f240331f5a80b27fb7442ca1ae5f81b5d0721e4143007ee0b4c4d1b1122c83c0e1aeaa5bea48167a6b2c3c92b3ad4d1100f7e768e607e
-  data.tar.gz: d7e0ea4e7ad204b9e5e6d11136008580a435bfe43e7ea9784926e6af60c0347dd556dd4d4410e542e29eb3c7a75be0fae3f6164a2bfbd7c451d37f121f62b932
+  metadata.gz: 2f646c1a4564960ba78494efef9ec64fa43eb9a4107d7c004884d977649d5cdf07aa19abec49b887d046d1c7a15ec7b61afbde60f4b12769c77f7578ba55df8e
+  data.tar.gz: b80f9fd21f045f396bceea9af860e471191669dc80b1d3e57ae38c5d1b7dcebe3fdfc1646774af043d6b8f243cd7335c9b8f044aff9232146e3c0a4dad8a58d8

data/CHANGELOG.md

@@ -1,5 +1,14 @@
 # CHANGELOG
 
+## v2.0.0.pre
+
+- Major change: RSpec plugin to use HTTP Upload API instead of websocket connection to send test data #174 #175 - @niceking
+- `identifier` field removed from trace #176 - @amybiyuliu
+
+## v1.5.0
+
+- Send `failure_expanded` from minitest #171 - @nprizal 
+
 ## v1.4.2
 
 - Update collector argument in the Analytics::API payload #170 - @KatieWright26

data/Gemfile.lock

@@ -1,23 +1,23 @@
 PATH
   remote: .
   specs:
-    buildkite-test_collector (1.4.2)
+    buildkite-test_collector (2.0.0.pre)
       activesupport (>= 4.2)
       websocket (~> 1.2)
 
 GEM
   remote: https://rubygems.org/
   specs:
-    activesupport (7.0.4.1)
+    activesupport (7.0.4.2)
       concurrent-ruby (~> 1.0, >= 1.0.2)
       i18n (>= 1.6, < 2)
       minitest (>= 5.1)
       tzinfo (~> 2.0)
-    concurrent-ruby (1.1.10)
+    concurrent-ruby (1.2.2)
     diff-lcs (1.4.4)
     i18n (1.12.0)
       concurrent-ruby (~> 1.0)
-    minitest (5.17.0)
+    minitest (5.18.0)
     rake (13.0.6)
     rspec (3.10.0)
       rspec-core (~> 3.10.0)
@@ -32,7 +32,7 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.10.0)
     rspec-support (3.10.3)
-    tzinfo (2.0.5)
+    tzinfo (2.0.6)
       concurrent-ruby (~> 1.0)
     websocket (1.2.9)
 
@@ -47,4 +47,4 @@ DEPENDENCIES
   rspec-expectations (~> 3.10)
 
 BUNDLED WITH
-   2.2.22
+   2.3.25

data/lib/buildkite/test_collector/http_client.rb

@@ -28,6 +28,28 @@ module Buildkite::TestCollector
       http.request(contact)
     end
 
+    def post_json(data)
+      contact_uri = URI.parse(url)
+
+      http = Net::HTTP.new(contact_uri.host, contact_uri.port)
+      http.use_ssl = contact_uri.scheme == "https"
+
+      contact = Net::HTTP::Post.new(contact_uri.path, {
+        "Authorization" => authorization_header,
+        "Content-Type" => "application/json",
+      })
+
+      data_set = data.map(&:as_hash)
+
+      contact.body = {
+        run_env: Buildkite::TestCollector::CI.env,
+        format: "json",
+        data: data_set
+      }.to_json
+
+      http.request(contact)
+    end
+
     private
 
     attr :url

data/lib/buildkite/test_collector/library_hooks/rspec.rb

@@ -11,8 +11,6 @@ Buildkite::TestCollector.uploader = Buildkite::TestCollector::Uploader
 RSpec.configure do |config|
   config.before(:suite) do
     config.add_formatter Buildkite::TestCollector::RSpecPlugin::Reporter
-
-    Buildkite::TestCollector.safe { Buildkite::TestCollector::Uploader.configure }
   end
 
   config.around(:each) do |example|

data/lib/buildkite/test_collector/minitest_plugin/trace.rb

@@ -15,12 +15,10 @@ module Buildkite::TestCollector::MinitestPlugin
 
     FILE_PATH_REGEX = /^(.*?\.(rb|feature))/
 
-    def initialize(example, history:, failure_reason: nil, failure_expanded: [])
+    def initialize(example, history:)
       @id = SecureRandom.uuid
       @example = example
       @history = history
-      @failure_reason = failure_reason
-      @failure_expanded = failure_expanded
     end
 
     def result
@@ -36,7 +34,6 @@ module Buildkite::TestCollector::MinitestPlugin
         id: id,
         scope: example.class.name,
         name: example.name,
-        identifier: identifier,
         location: location,
         file_name: file_name,
         result: result,
@@ -53,7 +50,6 @@ module Buildkite::TestCollector::MinitestPlugin
         "#{file_name}:#{line_number}"
       end
     end
-    alias_method :identifier, :location
 
     def file_name
       @file_name ||= File.join('./', source_location[0].delete_prefix(project_dir))
@@ -72,17 +68,20 @@ module Buildkite::TestCollector::MinitestPlugin
     end
 
     def failure_reason
-      @failure_reason ||= example.failure&.message
+      @failure_reason ||= strip_invalid_utf8_chars(example.failure&.message)&.split("\n")&.first
     end
 
     def failure_expanded
-      @failure_expanded ||= begin
-        example.failures.map do |failure|
-          {
-            expanded: failure.message,
-            backtrace: failure.backtrace,
-          }
-        end
+      @failure_expanded ||= example.failures.map.with_index do |failure, index|
+        # remove the first line of message from the first failure
+        # to avoid duplicate line in Test Analytics UI
+        messages = strip_invalid_utf8_chars(failure.message).split("\n")
+        messages = messages[1..] if index.zero?
+
+        {
+          expanded: messages,
+          backtrace: failure.backtrace
+        }
       end
     end
 

data/lib/buildkite/test_collector/rspec_plugin/reporter.rb

@@ -7,6 +7,7 @@ module Buildkite::TestCollector::RSpecPlugin
     attr_reader :output
 
     def initialize(output)
+      Buildkite::TestCollector.session = Buildkite::TestCollector::Session.new
       @output = output
     end
 
@@ -19,21 +20,13 @@ module Buildkite::TestCollector::RSpecPlugin
         if example.execution_result.status == :failed
           trace.failure_reason, trace.failure_expanded = failure_info(notification)
         end
-        Buildkite::TestCollector.session&.write_result(trace)
+        Buildkite::TestCollector.session.add_example_to_send_queue(example.id)
       end
     end
 
-    def dump_summary(notification)
-      if Buildkite::TestCollector.session.present?
-        examples_count = {
-          examples: notification.examples.count,
-          failed: notification.failed_examples.count,
-          pending: notification.pending_examples.count,
-          errors_outside_examples: notification.errors_outside_of_examples_count
-        }
-
-        Buildkite::TestCollector.session.close(examples_count)
-      end
+    def dump_summary(_notification)
+      Buildkite::TestCollector.session.send_remaining_data
+      Buildkite::TestCollector.session.close
     end
 
     alias_method :example_passed, :handle_example

data/lib/buildkite/test_collector/rspec_plugin/trace.rb

@@ -28,7 +28,6 @@ module Buildkite::TestCollector::RSpecPlugin
         id: id,
         scope: example.example_group.metadata[:full_description],
         name: example.description,
-        identifier: example.id,
         location: example.location,
         file_name: file_name,
         result: result,

data/lib/buildkite/test_collector/session.rb

@@ -2,332 +2,56 @@
 
 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
+    UPLOAD_THREAD_TIMEOUT = 60
+    UPLOAD_SESSION_TIMEOUT = 60
+    UPLOAD_API_MAX_RESULTS = 5000
 
-    # We keep a private reference so that mocking libraries won't break JSON
-    JSON_PARSE = JSON.method(:parse)
-    private_constant :JSON_PARSE
-
-    class RejectedSubscription < StandardError; end
-    class InitialConnectionFailure < StandardError; end
-
-    DISCONNECTED_EXCEPTIONS = [
-      Buildkite::TestCollector::SocketConnection::HandshakeError,
-      Buildkite::TestCollector::TimeoutError,
-      Buildkite::TestCollector::SocketConnection::SocketError,
-      RejectedSubscription,
-      InitialConnectionFailure,
-    ]
-
-    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 Buildkite::TestCollector::TimeoutError, InitialConnectionFailure => e
-        Buildkite::TestCollector.logger.warn("buildkite-test_collector 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 "buildkite-test_collector 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
+    def initialize
+      @send_queue_ids = []
+      @upload_threads = []
     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
+    def add_example_to_send_queue(id)
+      @send_queue_ids << id
 
-        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 "buildkite-test_collector 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
+      if @send_queue_ids.size >= Buildkite::TestCollector.batch_size
+        send_ids = @send_queue_ids.shift(Buildkite::TestCollector.batch_size)
+        upload_data(send_ids)
       end
-      retransmit
     end
 
-    def close(examples_count)
-      @closing = true
-      @examples_count = examples_count
-      Buildkite::TestCollector.logger.debug("closing socket connection")
+    def send_remaining_data
+      return if @send_queue_ids.empty?
 
-      # 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")
+      upload_data(@send_queue_ids)
     end
 
-    def handle(_connection, data)
-      data = JSON_PARSE.call(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 close
+      # There are two thread joins here, because the inner join will wait up to
+      # UPLOAD_THREAD_TIMEOUT seconds PER thread that is uploading data, i.e.
+      # n_threads x UPLOAD_THREAD_TIMEOUT latency if Buildkite happens to be
+      # down. By wrapping that in an outer thread join with the
+      # UPLOAD_SESSION_TIMEOUT, we ensure that we only wait a max of
+      # UPLOAD_SESSION_TIMEOUT seconds before the session exits.
+      Thread.new do
+        @upload_threads.each { |t| t.join(UPLOAD_THREAD_TIMEOUT) }
+      end.join(UPLOAD_SESSION_TIMEOUT)
 
-    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
+      @upload_threads.each { |t| t&.kill }
     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
+    def upload_data(ids)
+      data = Buildkite::TestCollector.uploader.traces.values_at(*ids).compact
 
-      @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
+      # we do this in batches of UPLOAD_API_MAX_RESULTS in case the number of
+      # results exceeds this due to a bug, or user error in configuring the
+      # batch size
+      data.each_slice(UPLOAD_API_MAX_RESULTS) do |batch|
+        new_thread = Buildkite::TestCollector::Uploader.upload(batch)
+        @upload_threads << new_thread if new_thread
       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

@@ -73,12 +73,15 @@ module Buildkite::TestCollector
           end
         end
       # These get re-raise from session, we should fail gracefully
-      rescue *Buildkite::TestCollector::Session::DISCONNECTED_EXCEPTIONS => e
+      rescue *Buildkite::TestCollector::SocketSession::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
+      rescue EOFError, OpenSSL::SSL::SSLError => e
+        # https://github.com/buildkite/test-collector-ruby/pull/147#issuecomment-1250485611
+        raise if e.class == OpenSSL::SSL::SSLError && e.message != "SSL_read: unexpected eof while reading"
+
         Buildkite::TestCollector.logger.warn("#{e}")
         if @socket
-          Buildkite::TestCollector.logger.error("attempting disconnected flow")
+          Buildkite::TestCollector.logger.warn("attempting disconnected flow")
           @session.disconnected(self)
           disconnect
         end

data/lib/buildkite/test_collector/socket_session.rb

@@ -0,0 +1,333 @@
+# frozen_string_literal: true
+
+module Buildkite::TestCollector
+  class SocketSession
+    # 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
+
+    # We keep a private reference so that mocking libraries won't break JSON
+    JSON_PARSE = JSON.method(:parse)
+    private_constant :JSON_PARSE
+
+    class RejectedSubscription < StandardError; end
+    class InitialConnectionFailure < StandardError; end
+
+    DISCONNECTED_EXCEPTIONS = [
+      Buildkite::TestCollector::SocketConnection::HandshakeError,
+      Buildkite::TestCollector::TimeoutError,
+      Buildkite::TestCollector::SocketConnection::SocketError,
+      RejectedSubscription,
+      InitialConnectionFailure,
+    ]
+
+    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 Buildkite::TestCollector::TimeoutError, InitialConnectionFailure => e
+        Buildkite::TestCollector.logger.warn("buildkite-test_collector 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 "buildkite-test_collector 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 "buildkite-test_collector 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.call(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/uploader.rb

@@ -2,6 +2,8 @@
 
 module Buildkite::TestCollector
   class Uploader
+    MAX_UPLOAD_ATTEMPTS = 3
+
     def self.traces
       @traces ||= {}
     end
@@ -17,6 +19,14 @@ module Buildkite::TestCollector
       EOFError
     ]
 
+    RETRYABLE_UPLOAD_ERRORS = [
+      Net::ReadTimeout,
+      Net::OpenTimeout,
+      OpenSSL::SSL::SSLError,
+      OpenSSL::SSL::SSLErrorWaitReadable,
+      EOFError
+    ]
+
     def self.configure
       Buildkite::TestCollector.logger.debug("hello from main thread")
 
@@ -38,7 +48,7 @@ module Buildkite::TestCollector
           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)
+            Buildkite::TestCollector.session = Buildkite::TestCollector::SocketSession.new(socket_url, http.authorization_header, channel)
           end
         else
           request_id = response.to_hash["x-request-id"]
@@ -54,5 +64,22 @@ module Buildkite::TestCollector
     def self.tracer
       Thread.current[:_buildkite_tracer]
     end
+
+    def self.upload(data)
+      return false unless Buildkite::TestCollector.api_token
+
+      http = Buildkite::TestCollector::HTTPClient.new(Buildkite::TestCollector.url)
+
+      Thread.new do
+        response = begin
+          upload_attempts ||= 0
+          http.post_json(data)
+        rescue *Buildkite::TestCollector::Uploader::RETRYABLE_UPLOAD_ERRORS => e
+          if (upload_attempts += 1) < MAX_UPLOAD_ATTEMPTS
+            retry
+          end
+        end
+      end
+    end
   end
 end

data/lib/buildkite/test_collector/version.rb

@@ -2,7 +2,7 @@
 
 module Buildkite
   module TestCollector
-    VERSION = "1.4.2"
+    VERSION = "2.0.0.pre"
     NAME = "buildkite-test_collector"
   end
 end

data/lib/buildkite/test_collector.rb

@@ -30,11 +30,13 @@ require_relative "test_collector/network"
 require_relative "test_collector/object"
 require_relative "test_collector/tracer"
 require_relative "test_collector/socket_connection"
+require_relative "test_collector/socket_session"
 require_relative "test_collector/session"
 
 module Buildkite
   module TestCollector
     DEFAULT_URL = "https://analytics-api.buildkite.com/v1/uploads"
+    DEFAULT_UPLOAD_BATCH_SIZE = 500
 
     class << self
       attr_accessor :api_token
@@ -45,6 +47,7 @@ module Buildkite
       attr_accessor :tracing_enabled
       attr_accessor :artifact_path
       attr_accessor :env
+      attr_accessor :batch_size
     end
 
     def self.configure(hook:, token: nil, url: nil, debug_enabled: false, tracing_enabled: true, artifact_path: nil, env: {})
@@ -54,7 +57,7 @@ module Buildkite
       self.tracing_enabled = tracing_enabled
       self.artifact_path = artifact_path
       self.env = env
-
+      self.batch_size = ENV.fetch("BUILDKITE_ANALYTICS_UPLOAD_BATCH_SIZE") { DEFAULT_UPLOAD_BATCH_SIZE }.to_i
       self.hook_into(hook)
     end
 

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: buildkite-test_collector
 version: !ruby/object:Gem::Version
-  version: 1.4.2
+  version: 2.0.0.pre
 platform: ruby
 authors:
 - Buildkite
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-01-19 00:00:00.000000000 Z
+date: 2023-03-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activesupport
@@ -104,6 +104,7 @@ files:
 - lib/buildkite/test_collector/rspec_plugin/trace.rb
 - lib/buildkite/test_collector/session.rb
 - lib/buildkite/test_collector/socket_connection.rb
+- lib/buildkite/test_collector/socket_session.rb
 - lib/buildkite/test_collector/tracer.rb
 - lib/buildkite/test_collector/uploader.rb
 - lib/buildkite/test_collector/version.rb
@@ -125,11 +126,11 @@ required_ruby_version: !ruby/object:Gem::Requirement
       version: 2.3.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
-  - - ">="
+  - - ">"
     - !ruby/object:Gem::Version
-      version: '0'
+      version: 1.3.1
 requirements: []
-rubygems_version: 3.3.26
+rubygems_version: 3.1.6
 signing_key: 
 specification_version: 4
 summary: Track test executions and report to Buildkite Test Analytics