Asynchronously load next records in ActiveRecordCursor

lib/job-iteration/active_record_cursor.rb

@@ -7,6 +7,7 @@ class ActiveRecordCursor # @private
     include Comparable
 
     attr_reader :position
+    attr_reader :next_relation
     attr_accessor :reached_end
 
     class ConditionNotSupportedError < ArgumentError
@@ -33,6 +34,7 @@ def initialize(relation, columns = nil, position = nil)
 
       @base_relation = relation.reorder(@columns.join(","))
       @reached_end = false
+      @async_enabled = relation.connection.try(:async_enabled?)
     end
 
     def <=>(other)
@@ -59,11 +61,7 @@ def update_from_record(record)
     def next_batch(batch_size)
       return nil if @reached_end
 
-      relation = @base_relation.limit(batch_size)
-
-      if (conditions = self.conditions).any?
-        relation = relation.where(*conditions)
-      end
+      relation = (@next_relation || build_relation(batch_size))
 
       records = relation.uncached do
         relation.to_a
@@ -72,7 +70,12 @@ def next_batch(batch_size)
       update_from_record(records.last) unless records.empty?
       @reached_end = records.size < batch_size
 
-      records.empty? ? nil : records
+      if records.empty?
+        nil
+      else
+        @next_relation = async_preload(batch_size)
+        records
+      end
     end
 
     protected
@@ -94,5 +97,24 @@ def conditions
       ret.pop
       ret
     end
+
+    private
+
+    def build_relation(batch_size)
+      relation = @base_relation.limit(batch_size)
+
+      if (conditions = self.conditions).any?
+        relation = relation.where(*conditions)
+      end
+
+      relation
+    end
+
+    def async_preload(batch_size)
+      return unless @async_enabled
+
+      relation = build_relation(batch_size)
+      relation.uncached { relation.load_async }
+    end
   end
 end

test/test_helper.rb

@@ -54,6 +54,9 @@ class Product < ActiveRecord::Base
 }
 connection_config[:password] = "root" if ENV["CI"]
 
+if ActiveRecord.respond_to?(:async_query_executor)
+  ActiveRecord.async_query_executor = :global_thread_pool
+end
 ActiveRecord::Base.establish_connection(connection_config)
 
 Redis.singleton_class.class_eval do

test/unit/active_record_cursor_test.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+module JobIteration
+  class ActiveRecordCursorTest < IterationUnitTest
+    test "#next_batch preloads the following batch asynchronously" do
+      skip("Only supported on Rails >= 7") unless Product.connection.try(:async_enabled?)
+
+      cursor = ActiveRecordCursor.new(Product.all)
+      cursor.next_batch(2)
+
+      assert_predicate(cursor.next_relation, :scheduled?)
+    end
+  end
+end