Fix jobs being added to batch after they might already execute (#35496)

app/models/worker_batch.rb

@@ -19,17 +19,22 @@ class WorkerBatch
     redis.hset(key, { 'async_refresh_key' => async_refresh_key, 'threshold' => threshold })
   end
 
+  def within
+    raise NoBlockGivenError unless block_given?
+
+    begin
+      Thread.current[:batch] = self
+      yield
+    ensure
+      Thread.current[:batch] = nil
+    end
+  end
+
   # Add jobs to the batch. Usually when the batch is created.
   # @param [Array<String>] jids
   def add_jobs(jids)
     if jids.blank?
-      async_refresh_key = redis.hget(key, 'async_refresh_key')
+      finish!
-
-      if async_refresh_key.present?
-        async_refresh = AsyncRefresh.new(async_refresh_key)
-        async_refresh.finish!
-      end
-
       return
     end
 
@@ -55,8 +60,23 @@ class WorkerBatch
     if async_refresh_key.present?
       async_refresh = AsyncRefresh.new(async_refresh_key)
       async_refresh.increment_result_count(by: 1)
-      async_refresh.finish! if pending.zero? || processed >= threshold.to_f * (processed + pending)
     end
+
+    if pending.zero? || processed >= (threshold || 1.0).to_f * (processed + pending)
+      async_refresh&.finish!
+      cleanup
+    end
+  end
+
+  def finish!
+    async_refresh_key = redis.hget(key, 'async_refresh_key')
+
+    if async_refresh_key.present?
+      async_refresh = AsyncRefresh.new(async_refresh_key)
+      async_refresh.finish!
+    end
+
+    cleanup
   end
 
   # Get pending jobs.
@@ -76,4 +96,8 @@ class WorkerBatch
   def key(suffix = nil)
     "worker_batch:#{@id}#{":#{suffix}" if suffix}"
   end
+
+  def cleanup
+    redis.del(key, key('jobs'))
+  end
 end

app/services/activitypub/fetch_replies_service.rb

@@ -17,7 +17,12 @@ class ActivityPub::FetchRepliesService < BaseService
 
     batch = WorkerBatch.new
     batch.connect(async_refresh_key) if async_refresh_key.present?
-    batch.add_jobs(FetchReplyWorker.push_bulk(@items) { |reply_uri| [reply_uri, { 'request_id' => request_id, 'batch_id' => batch.id }] })
+    batch.finish! if @items.empty?
+    batch.within do
+      FetchReplyWorker.push_bulk(@items) do |reply_uri|
+        [reply_uri, { 'request_id' => request_id, 'batch_id' => batch.id }]
+      end
+    end
 
     [@items, n_pages]
   end

config/initializers/sidekiq.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative '../../lib/mastodon/sidekiq_middleware'
+require_relative '../../lib/mastodon/worker_batch_middleware'
 
 Sidekiq.configure_server do |config|
   config.redis = REDIS_CONFIGURATION.sidekiq
@@ -72,14 +73,12 @@ Sidekiq.configure_server do |config|
 
   config.server_middleware do |chain|
     chain.add Mastodon::SidekiqMiddleware
-  end
-
-  config.server_middleware do |chain|
     chain.add SidekiqUniqueJobs::Middleware::Server
   end
 
   config.client_middleware do |chain|
     chain.add SidekiqUniqueJobs::Middleware::Client
+    chain.add Mastodon::WorkerBatchMiddleware
   end
 
   config.on(:startup) do
@@ -105,6 +104,7 @@ Sidekiq.configure_client do |config|
 
   config.client_middleware do |chain|
     chain.add SidekiqUniqueJobs::Middleware::Client
+    chain.add Mastodon::WorkerBatchMiddleware
   end
 
   config.logger.level = Logger.const_get(ENV.fetch('RAILS_LOG_LEVEL', 'info').upcase.to_s)

lib/mastodon/worker_batch_middleware.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+class Mastodon::WorkerBatchMiddleware
+  def call(_worker, msg, _queue, _redis_pool = nil)
+    if (batch = Thread.current[:batch])
+      batch.add_jobs([msg['jid']])
+    end
+
+    yield
+  end
+end