Merge pull request #323 from Shopify/ack-race-cond

Refactor test ACK to reduce the race condition window

Author: ChrisBr

ruby/lib/ci/queue/redis/build_record.rb

@@ -56,20 +56,29 @@ def record_warning(type, attributes)
           redis.rpush(key('warnings'), Marshal.dump([type, attributes]))
         end
 
+        Test = Struct.new(:id) # Hack
+
         def record_error(id, payload, stats: nil)
-          redis.pipelined do |pipeline|
-            pipeline.hset(
-              key('error-reports'),
-              id.dup.force_encoding(Encoding::BINARY),
-              payload.dup.force_encoding(Encoding::BINARY),
-            )
-            pipeline.expire(key('error-reports'), config.redis_ttl)
-            record_stats(stats, pipeline: pipeline)
+          # FIXME: the ack and hset should be atomic
+          # otherwise the reporter may see the empty queue before the error
+          # is appended and wrongly think it's a success.
+          if @queue.acknowledge(Test.new(id))
+            redis.pipelined do |pipeline|
+              pipeline.hset(
+                key('error-reports'),
+                id.dup.force_encoding(Encoding::BINARY),
+                payload.dup.force_encoding(Encoding::BINARY),
+              )
+              pipeline.expire(key('error-reports'), config.redis_ttl)
+              record_stats(stats, pipeline: pipeline)
+              @queue.increment_test_failed
+            end
           end
           nil
         end
 
-        def record_success(id, stats: nil, skip_flaky_record: false)
+        def record_success(id, stats: nil, skip_flaky_record: false, acknowledge: true)
+          @queue.acknowledge(Test.new(id)) if acknowledge
           error_reports_deleted_count, requeued_count, _ = redis.pipelined do |pipeline|
             pipeline.hdel(key('error-reports'), id.dup.force_encoding(Encoding::BINARY))
             pipeline.hget(key('requeues-count'), id.b)

ruby/lib/minitest/queue.rb

@@ -150,6 +150,7 @@ def self.relative_path(path, root: project_root)
     end
 
     class SingleExample
+      attr_reader :method_name
 
       def initialize(runnable, method_name)
         @runnable = runnable
@@ -250,15 +251,8 @@ def run_from_queue(reporter, *)
 
         if failed && CI::Queue.requeueable?(result) && queue.requeue(example)
           result.requeue!
-          reporter.record(result)
-        elsif queue.acknowledge(example)
-          reporter.record(result)
-          queue.increment_test_failed if failed
-        elsif !failed
-          # If the test was already acknowledged by another worker (we timed out)
-          # Then we only record it if it is successful.
-          reporter.record(result)
         end
+        reporter.record(result)
       end
       queue.stop_heartbeat!
     rescue Errno::EPIPE

ruby/lib/minitest/queue/build_status_recorder.rb

@@ -52,7 +52,7 @@ def record(test)
         if (test.failure || test.error?) && !test.skipped?
           build.record_error("#{test.klass}##{test.name}", dump(test), stats: stats)
         else
-          build.record_success("#{test.klass}##{test.name}", stats: stats, skip_flaky_record: test.skipped?)
+          build.record_success("#{test.klass}##{test.name}", stats: stats, skip_flaky_record: test.skipped?, acknowledge: !test.requeued?)
         end
       end
 

ruby/lib/rspec/queue.rb

@@ -224,13 +224,8 @@ def finish(reporter, acknowledge: true)
             reporter.cancel_run!
             dup.mark_as_requeued!(reporter)
             return true
-          elsif reporter.acknowledge || !@exception
-            # If the test was already acknowledged by another worker (we timed out)
-            # Then we only record it if it is successful.
-            super(reporter)
           else
-            reporter.cancel_run!
-            return
+            super(reporter)
           end
         else
           super(reporter)

ruby/test/integration/minitest_redis_test.rb

@@ -514,7 +514,7 @@ def test_retry_report
           'skips' => 0,
           'requeues' => 0,
           'total_time' => index + 1,
-        })
+        }, acknowledge: false)
       end
 
       # Retry first worker, bailing out

ruby/test/minitest/queue/build_status_recorder_test.rb

@@ -15,19 +15,28 @@ def setup
     end
 
     def test_aggregation
+      reserve(@queue, "a")
       @reporter.record(result('a', failure: "Something went wrong"))
+      reserve(@queue, "b")
       @reporter.record(result('b', unexpected_error: true))
+      reserve(@queue, "h")
       @reporter.record(result('h', failure: "Something went wrong", requeued: true))
 
       second_queue = worker(2)
       second_reporter = BuildStatusRecorder.new(build: second_queue.build)
       second_reporter.start
 
+      reserve(second_queue, "c")
       second_reporter.record(result('c', failure: "Something went wrong"))
+      reserve(second_queue, "d")
       second_reporter.record(result('d', unexpected_error: true))
+      reserve(second_queue, "e")
       second_reporter.record(result('e', skipped: true))
+      reserve(second_queue, "f")
       second_reporter.record(result('f', unexpected_error: true))
+      reserve(second_queue, "g")
       second_reporter.record(result('g', requeued: true))
+      reserve(second_queue, "h")
       second_reporter.record(result('h', skipped: true, requeued: true))
 
       assert_equal 9, summary.assertions
@@ -40,19 +49,39 @@ def test_aggregation
     end
 
     def test_retrying_test
-      @reporter.record(result('a', failure: "Something went wrong"))
-      assert_equal 1, summary.error_reports.size
+      yielded = false
+
+      test = nil
+
+      @queue.poll do |_test|
+        test = _test
+        assert_equal "a", test.method_name
+        @reporter.record(result(test.method_name, failure: "Something went wrong"))
+
+        assert_equal 1, summary.error_reports.size
+
+        yielded = true
+        break
+      end
+
+      assert yielded, "@queue.poll didn't yield"
 
       second_queue = worker(2)
       second_reporter = BuildStatusRecorder.new(build: second_queue.build)
       second_reporter.start
 
-      second_reporter.record(result('a'))
+      # pretend we reserved the same test again
+      reserve(second_queue, "a")
+      second_reporter.record(result("a"))
       assert_equal 0, summary.error_reports.size
     end
 
     private
 
+    def reserve(queue, method_name)
+      queue.instance_variable_set(:@reserved_test, Minitest::Queue::SingleExample.new("Minitest::Test", method_name).id)
+    end
+
     def worker(id)
       CI::Queue::Redis.new(
         @redis_url,
@@ -61,7 +90,9 @@ def worker(id)
           worker_id: id.to_s,
           timeout: 0.2,
         ),
-      ).populate([])
+      ).populate([
+        Minitest::Queue::SingleExample.new("Minitest::Test", "a")
+      ])
     end
 
     def summary