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')
-
-      if async_refresh_key.present?
-        async_refresh = AsyncRefresh.new(async_refresh_key)
-        async_refresh.finish!
-      end
-
+      finish!
       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