From 64491d6aa464e80338d21778444f6bf34545c405 Mon Sep 17 00:00:00 2001
From: Elson Oliveira <elson.oliveira@shopify.com>
Date: Tue, 21 Feb 2023 14:40:47 -0700
Subject: [PATCH] Add active record cursor next iteration metric

---
 lib/job-iteration/active_record_enumerator.rb |  8 ++++-
 test/unit/active_record_enumerator_test.rb    | 32 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/lib/job-iteration/active_record_enumerator.rb b/lib/job-iteration/active_record_enumerator.rb
index 363a4ecf..3aa37839 100644
--- a/lib/job-iteration/active_record_enumerator.rb
+++ b/lib/job-iteration/active_record_enumerator.rb
@@ -31,7 +31,7 @@ def records
     def batches
       cursor = finder_cursor
       Enumerator.new(method(:size)) do |yielder|
-        while (records = cursor.next_batch(@batch_size))
+        while (records = instrument_next_batch(cursor))
           yielder.yield(records, cursor_value(records.last)) if records.any?
         end
       end
@@ -43,6 +43,12 @@ def size
 
     private
 
+    def instrument_next_batch(cursor)
+      ActiveSupport::Notifications.instrument("active_record_cursor.iteration") do
+        cursor.next_batch(@batch_size)
+      end
+    end
+
     def cursor_value(record)
       positions = @columns.map do |column|
         attribute_name = column.to_s.split(".").last
diff --git a/test/unit/active_record_enumerator_test.rb b/test/unit/active_record_enumerator_test.rb
index f1a83ab1..c6955718 100644
--- a/test/unit/active_record_enumerator_test.rb
+++ b/test/unit/active_record_enumerator_test.rb
@@ -57,6 +57,38 @@ class ActiveRecordEnumeratorTest < IterationUnitTest
       assert_equal([shops, shops.last.id], enum.first)
     end
 
+    class StubbedCursor
+      include ActiveSupport::Testing::TimeHelpers
+      def initialize(wait_time:)
+        @wait_time = wait_time
+      end
+
+      def next_batch(*)
+        travel(@wait_time)
+      end
+    end
+
+    test "enumerator next batch is instrumented with proper duration" do
+      wait_time = 15.seconds
+      freeze_time do
+        stubbed_cursor = StubbedCursor.new(wait_time: wait_time)
+
+        ActiveSupport::Notifications.subscribe("active_record_cursor.iteration") do |*args|
+          ActiveSupport::Notifications.unsubscribe("active_record_cursor.iteration")
+          event = ActiveSupport::Notifications::Event.new(*args)
+          assert_equal(wait_time.in_milliseconds, event.duration)
+        end
+        enum = build_enumerator
+        enum.send(:instrument_next_batch, stubbed_cursor)
+      end
+    end
+
+    test "enumerator next batch is instrumented" do
+      ActiveSupport::Notifications.expects(:instrument).with("active_record_cursor.iteration")
+      enum = build_enumerator.batches
+      enum.first
+    end
+
     test "columns are configurable" do
       enum = build_enumerator(columns: [:updated_at]).batches
       shops = Product.order(:updated_at).take(2)