From 02679f5981607569918025dc1654fd9d7482ad87 Mon Sep 17 00:00:00 2001
From: Tessa Bradbury <tessa@buildkite.com>
Date: Wed, 11 Mar 2026 16:37:38 +1100
Subject: [PATCH] Add support for instance-level throttles

---
 README.md                                  | 22 ++++++++++++++++
 app/models/maintenance_tasks/task.rb       | 29 +++++++++++++++++++++-
 test/models/maintenance_tasks/task_test.rb | 14 +++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 109d32b7..8b94cc20 100644
--- a/README.md
+++ b/README.md
@@ -467,6 +467,28 @@ module Maintenance
 end
 ```
 
+If you want the throttle to change based on task attributes, you can define the
+`throttle_on` inside `#initialize`:
+
+```ruby
+# app/tasks/maintenance/update_posts_throttled_task.rb
+
+module Maintenance
+  class UpdatePostsThrottledTask < MaintenanceTasks::Task
+    attribute :throttle_backoff_seconds, :integer
+
+    def initialize(*)
+      super
+
+      throttle_on(backoff: throttle_backoff_seconds) do
+        DatabaseStatus.unhealthy?
+      end
+    end
+    # ...
+  end
+end
+```
+
 ### Custom Task Parameters
 
 Tasks may need additional information, supplied via parameters, to run.
diff --git a/app/models/maintenance_tasks/task.rb b/app/models/maintenance_tasks/task.rb
index c089ea1e..d4743b47 100644
--- a/app/models/maintenance_tasks/task.rb
+++ b/app/models/maintenance_tasks/task.rb
@@ -17,7 +17,7 @@ class NotFoundError < NameError; end
     # backoff. Note that Tasks inherit conditions from their superclasses.
     #
     # @api private
-    class_attribute :throttle_conditions, default: []
+    class_attribute :throttle_conditions, default: [], instance_reader: false
 
     # The number of active records to fetch in a single query when iterating
     # over an Active Record collection task.
@@ -331,5 +331,32 @@ def count
     def enumerator_builder(cursor:)
       nil
     end
+
+    # Add a condition under which this Task instance will be throttled.
+    #
+    # @param backoff [ActiveSupport::Duration, #call] a custom backoff
+    #   can be specified. This is the time to wait before retrying the Task,
+    #   defaulting to 30 seconds. If provided as a Duration, the backoff is
+    #   wrapped in a proc. Alternatively,an object responding to call can be
+    #   used. It must return an ActiveSupport::Duration.
+    # @yieldreturn [Boolean] where the throttle condition is being met,
+    #   indicating that the Task should throttle.
+    def throttle_on(backoff: 30.seconds, &condition)
+      backoff_as_proc = backoff
+      backoff_as_proc = -> { backoff } unless backoff.respond_to?(:call)
+
+      @throttle_conditions ||= []
+      @throttle_conditions += [{ throttle_on: condition, backoff: backoff_as_proc }]
+    end
+
+    # The throttle conditions for a given Task instance. This is provided as
+    # an array of hashes, with each hash specifying two keys:
+    # throttle_condition and backoff. Note that Tasks inherit conditions from
+    # their superclasses.
+    #
+    # @api private
+    def throttle_conditions
+      self.class.throttle_conditions + (@throttle_conditions || [])
+    end
   end
 end
diff --git a/test/models/maintenance_tasks/task_test.rb b/test/models/maintenance_tasks/task_test.rb
index 572e69d5..57ba54fa 100644
--- a/test/models/maintenance_tasks/task_test.rb
+++ b/test/models/maintenance_tasks/task_test.rb
@@ -113,6 +113,20 @@ class TaskTest < ActiveSupport::TestCase
       Maintenance::TestTask.throttle_conditions = []
     end
 
+    test "#throttle_on registers throttle condition for Task" do
+      throttle_condition = -> { true }
+
+      task = Maintenance::TestTask.new
+      task.throttle_on(&throttle_condition)
+
+      task_throttle_conditions = task.throttle_conditions
+      assert_equal(1, task_throttle_conditions.size)
+
+      condition = task_throttle_conditions.first
+      assert_equal(throttle_condition, condition[:throttle_on])
+      assert_equal(30.seconds, condition[:backoff].call)
+    end
+
     test ".cursor_columns returns nil" do
       task = Task.new
       assert_nil task.cursor_columns