Reclaim unscheduled timer tasks instead of holding slots until run_time

src/bthread/timer_thread.cpp

@@ -34,6 +34,31 @@ namespace bthread {
 
 DEFINE_uint32(brpc_timer_num_buckets, 13, "brpc timer num buckets");
 
+// Tasks unscheduled after being pulled into the timer thread's min-heap are
+// only recycled when popped at their run_time, which can be far in the future
+// for large timeouts. To bound the memory they occupy (~ qps * timeout), the
+// timer thread periodically sweeps the heap and drops unscheduled tasks. The
+// sweep only kicks in once the heap grows beyond this size, so that small
+// heaps (where the retained memory is negligible) never pay the O(N) cost.
+DEFINE_uint32(brpc_timer_heap_sweep_min_size, 4096,
+              "The timer thread sweeps unscheduled tasks out of its internal "
+              "heap only when the heap has at least this many tasks");
+
+// The timer thread only consumes buckets and reclaims unscheduled tasks when
+// it wakes up, which normally happens at the nearest task's run_time. If every
+// pending task has a far-future run_time (e.g. minutes away), the thread would
+// sleep that whole time while newly scheduled-then-unscheduled tasks pile up
+// in the buckets, occupying pooled slots for the entire duration. Capping the
+// sleep makes the thread wake up periodically to drain the buckets and sweep
+// the heap, bounding that latency regardless of the timeout distribution.
+// 0 (the default) disables the cap: sleep until the nearest run_time, the
+// legacy behavior. Set it to a positive value to bound reclaim latency when
+// tasks may have far-future run_times.
+DEFINE_uint32(brpc_timer_max_wakeup_interval_ms, 0,
+              "The timer thread wakes up at least this often (in milliseconds) "
+              "to reclaim unscheduled tasks even when all pending tasks are far "
+              "in the future; 0 means no periodic wakeup");
+
 // Defined in task_control.cpp
 void run_worker_startfn();
 
@@ -132,6 +157,7 @@ TimerThread::TimerThread()
     , _buckets(NULL)
     , _nearest_run_time(std::numeric_limits<int64_t>::max())
     , _nsignals(0)
+    , _npending(0)
     , _thread(0) {
 }
 
@@ -327,6 +353,11 @@ void TimerThread::run() {
     // min heap of tasks (ordered by run_time)
     std::vector<Task*> tasks;
     tasks.reserve(4096);
+    // Heap size at the last sweep. Used to trigger the next sweep only after
+    // the heap has roughly doubled, keeping the amortized cost of sweeping at
+    // O(1) per task. It also follows the heap down when tasks are run, so a
+    // regrowth (e.g. filled with newly-unscheduled tasks) triggers a sweep.
+    size_t last_sweep_size = 0;
 
     // vars
     size_t nscheduled = 0;
@@ -374,6 +405,29 @@ void TimerThread::run() {
             }
         }
 
+        // A task is only checked by try_delete() once, right when it is pulled
+        // out of its bucket. If it gets unscheduled afterwards, it lingers in
+        // the heap (occupying a pooled Task slot) until it is popped at its
+        // run_time. For large timeouts this keeps a lot of dead tasks around.
+        // Sweep them out here. The sweep is gated on the heap having grown to
+        // twice its post-sweep size (and past a minimum), so the O(N) pass is
+        // amortized O(1) per task and small heaps never pay for it.
+        if (tasks.size() >= FLAGS_brpc_timer_heap_sweep_min_size &&
+            tasks.size() >= last_sweep_size * 2) {
+            size_t j = 0;
+            for (size_t i = 0; i < tasks.size(); ++i) {
+                Task* task = tasks[i];
+                if (!task->try_delete()) {  // still scheduled, keep it
+                    tasks[j++] = task;
+                }
+            }
+            if (j != tasks.size()) {
+                tasks.resize(j);
+                std::make_heap(tasks.begin(), tasks.end(), task_greater);
+            }
+            last_sweep_size = tasks.size();
+        }
+
         bool pull_again = false;
         while (!tasks.empty()) {
             Task* task1 = tasks[0];  // the about-to-run task
@@ -404,10 +458,18 @@ void TimerThread::run() {
                 ++ntriggered;
             }
         }
+        // Publish the heap size before possibly looping back on pull_again,
+        // so the observability counter doesn't go stale during the retry spin.
+        _npending.store((int64_t)tasks.size(), butil::memory_order_relaxed);
         if (pull_again) {
             BT_VLOG << "pull again, tasks=" << tasks.size();
             continue;
         }
+        // Let the sweep baseline follow the heap down as tasks are run, so
+        // that a heap refilled with (soon unscheduled) tasks is swept again.
+        if (tasks.size() < last_sweep_size) {
+            last_sweep_size = tasks.size();
+        }
 
         // The realtime to wait for.
         int64_t next_run_time = std::numeric_limits<int64_t>::max();
@@ -434,7 +496,18 @@ void TimerThread::run() {
         timespec next_timeout = { 0, 0 };
         const int64_t now = butil::gettimeofday_us();
         if (next_run_time != std::numeric_limits<int64_t>::max()) {
-            next_timeout = butil::microseconds_to_timespec(next_run_time - now);
+            int64_t wait_us = next_run_time - now;
+            // Cap the sleep so we periodically wake up to drain buckets and
+            // sweep the heap even when the nearest task is far in the future.
+            // Note: an empty heap keeps ptimeout NULL (sleep until woken by a
+            // schedule()), which is safe because the first task after the heap
+            // empties is always earlier than _nearest_run_time and wakes us.
+            const int64_t max_wakeup_us =
+                (int64_t)FLAGS_brpc_timer_max_wakeup_interval_ms * 1000;
+            if (max_wakeup_us > 0 && wait_us > max_wakeup_us) {
+                wait_us = max_wakeup_us;
+            }
+            next_timeout = butil::microseconds_to_timespec(wait_us);
             ptimeout = &next_timeout;
         }
         busy_seconds += (now - last_sleep_time) / 1000000.0;

src/bthread/timer_thread.h

@@ -84,7 +84,7 @@ class TimerThread {
     // Get identifier of internal pthread.
     // Returns (pthread_t)0 if start() is not called yet.
     pthread_t thread_id() const { return _thread; }
-    
+
 private:
     // the timer thread will run this method.
     void run();
@@ -100,6 +100,10 @@ class TimerThread {
     // the futex for wake up timer thread. can't use _nearest_run_time because
     // it's 64-bit.
     int _nsignals;
+    // Number of tasks buffered in the internal min-heap, published by the
+    // timer thread each iteration. Not part of the public API; read by unit
+    // tests through the -fno-access-control build flag.
+    butil::atomic<int64_t> _npending;
     pthread_t _thread;       // all scheduled task will be run on this thread
 };
 

test/bthread_timer_thread_unittest.cpp

@@ -17,6 +17,9 @@
 
 #include <gtest/gtest.h>
 #include <gflags/gflags.h>
+#include <algorithm>
+#include <string>
+#include <vector>
 #include "bthread/sys_futex.h"
 #include "bthread/timer_thread.h"
 #include "bthread/bthread.h"
@@ -254,4 +257,125 @@ TEST(TimerThreadTest, schedule_and_unschedule_in_task) {
     keeper5.expect_first_run();
 }
 
+static void noop_routine(void*) {}
+
+// RAII helper: set a gflag for the duration of a test and restore its previous
+// value on scope exit, so tests don't leak configuration into later tests and
+// cause order-dependent failures. Also asserts the flag exists and was set.
+class ScopedFlag {
+public:
+    ScopedFlag(const char* name, const char* value) : _name(name) {
+        EXPECT_TRUE(GFLAGS_NAMESPACE::GetCommandLineOption(name, &_old_value))
+            << "unknown gflag " << name;
+        EXPECT_FALSE(
+            GFLAGS_NAMESPACE::SetCommandLineOption(name, value).empty())
+            << "failed to set gflag " << name << "=" << value;
+    }
+    ~ScopedFlag() {
+        GFLAGS_NAMESPACE::SetCommandLineOption(_name.c_str(), _old_value.c_str());
+    }
+private:
+    std::string _name;
+    std::string _old_value;
+};
+
+// Tasks that are unscheduled after being pulled into the timer thread's
+// internal heap must not stay there (occupying a pooled Task slot) until
+// their run_time. With large timeouts that would let memory grow ~ qps *
+// timeout. The timer thread should sweep them out and keep the heap bounded.
+TEST(TimerThreadTest, sweep_unscheduled_tasks_in_heap) {
+    // Lower the sweep threshold so the test doesn't need a huge heap.
+    ScopedFlag sweep_flag("brpc_timer_heap_sweep_min_size", "512");
+
+    bthread::TimerThread timer_thread;
+    ASSERT_EQ(0, timer_thread.start(NULL));
+
+    // Run far enough in the future that these tasks never fire on their own.
+    const timespec far = butil::seconds_from_now(100000);
+    const size_t kBatch = 2000;
+    const size_t kRounds = 20;
+
+    int64_t max_pending = 0;
+    for (size_t r = 0; r < kRounds; ++r) {
+        std::vector<bthread::TimerThread::TaskId> ids;
+        ids.reserve(kBatch);
+        for (size_t i = 0; i < kBatch; ++i) {
+            ids.push_back(timer_thread.schedule(noop_routine, NULL, far));
+        }
+        // A near-term task forces the timer thread to wake up and consume the
+        // buckets, so the far tasks above land in the heap (alive).
+        timer_thread.schedule(noop_routine, NULL,
+                              butil::milliseconds_from_now(1));
+        usleep(20000);  // let the timer thread consume the buckets
+
+        // Now unschedule the far tasks: they become dead-in-heap, exactly the
+        // case that used to linger until run_time.
+        for (size_t i = 0; i < ids.size(); ++i) {
+            timer_thread.unschedule(ids[i]);
+        }
+        // Another near-term task wakes the timer thread again, triggering the
+        // sweep that reclaims the dead tasks.
+        timer_thread.schedule(noop_routine, NULL,
+                              butil::milliseconds_from_now(1));
+        usleep(20000);
+
+        // Read the internal heap size directly (brpc tests are built with
+        // -fno-access-control, so no public accessor is needed).
+        const int64_t pending =
+            timer_thread._npending.load(butil::memory_order_relaxed);
+        LOG(INFO) << "round=" << r << " pending=" << pending;
+        max_pending = std::max(max_pending, pending);
+    }
+
+    // Without the sweep, pending would grow to ~ kBatch * kRounds (40000).
+    // With it, the heap is bounded by roughly a couple of sweep thresholds.
+    EXPECT_LT(max_pending, (int64_t)(kBatch * kRounds) / 4)
+        << "dead tasks are not being reclaimed from the heap";
+
+    timer_thread.stop_and_join();
+}
+
+// When every pending task is far in the future, the nearest run_time never
+// arrives, so the timer thread used to sleep for the whole duration and never
+// consume the buckets. Tasks scheduled meanwhile (and the slots they occupy)
+// would then be stranded in the buckets until that far deadline. The capped
+// wakeup makes the timer thread wake up periodically to drain the buckets so
+// their tasks are consumed (and, if unscheduled, reclaimed) within the cap
+// rather than after the timeout.
+TEST(TimerThreadTest, periodic_wakeup_drains_buckets) {
+    ScopedFlag wakeup_flag("brpc_timer_max_wakeup_interval_ms", "50");
+
+    bthread::TimerThread timer_thread;
+    ASSERT_EQ(0, timer_thread.start(NULL));
+
+    // Anchor task an hour out: it becomes the nearest task, so the tasks below
+    // (with even later run_times) are never the "earliest" and thus never wake
+    // the timer via schedule() -- only the periodic wakeup can drain them.
+    timer_thread.schedule(noop_routine, NULL, butil::seconds_from_now(3600));
+    usleep(100000);  // let the anchor be consumed into the heap
+    // Only the anchor is in the heap so far.
+    ASSERT_EQ(1, timer_thread._npending.load(butil::memory_order_relaxed));
+
+    // Pile far tasks with strictly-increasing run_times into the buckets. None
+    // of these wake the timer.
+    const int kN = 2000;
+    for (int i = 0; i < kN; ++i) {
+        timer_thread.schedule(noop_routine, NULL,
+                              butil::seconds_from_now(3600 + 1 + i));
+    }
+
+    // Without the periodic wakeup the timer would stay asleep (nearest task is
+    // an hour away) and these would sit in the buckets, unconsumed. With it,
+    // they are pulled into the heap within a few wakeup intervals.
+    usleep(300000);  // several 50ms intervals
+    const int64_t pending =
+        timer_thread._npending.load(butil::memory_order_relaxed);
+    LOG(INFO) << "pending after bucket fill = " << pending << " (scheduled "
+              << kN << " + 1 anchor)";
+    EXPECT_GE(pending, (int64_t)kN)
+        << "buckets were not drained by the periodic wakeup";
+
+    timer_thread.stop_and_join();
+}
+
 } // end namespace