[Feature] Support Idle Callback in TaskGroup for flexibility usage

- Implement Idle Hook: Added TaskGroup::SetWorkerIdleCallback to allow executing custom logic (e.g., IO polling) when a worker thread is idle.
- Support Timeout Wait: Modified ParkingLot::wait to support an optional timeout, preventing workers from sleeping indefinitely when an idle callback is registered.
- Enable Thread-per-Core IO: Enabled thread-local IO management (like io_uring ) by invoking the hook within the worker's thread context.
- Add Unit Test: Added bthread_idle_unittest to verify worker isolation and idle callback execution.

Author: Kuankuan Guo

src/bthread/bthread.cpp

@@ -90,6 +90,8 @@ extern BAIDU_THREAD_LOCAL TaskGroup* tls_task_group;
 EXTERN_BAIDU_VOLATILE_THREAD_LOCAL(TaskGroup*, tls_task_group);
 extern void (*g_worker_startfn)();
 extern void (*g_tagged_worker_startfn)(bthread_tag_t);
+extern bool (*g_worker_idle_fn)(void);
+extern timespec g_worker_idle_timeout;
 extern void* (*g_create_span_func)();
 
 inline TaskControl* get_task_control() {
@@ -597,6 +599,24 @@ int bthread_set_tagged_worker_startfn(void (*start_fn)(bthread_tag_t)) {
     return 0;
 }
 
+int bthread_set_worker_idle_callback(bool (*fn)(void), uint64_t timeout_us) {
+    // Allow clearing the callback by passing NULL
+    if (fn == NULL) {
+        bthread::g_worker_idle_fn = NULL;
+        bthread::g_worker_idle_timeout = {0, 0};
+        return 0;
+    }
+    if (timeout_us == 0) {
+        return EINVAL;
+    }
+    timespec wait_time;
+    wait_time.tv_sec = timeout_us / 1000000;
+    wait_time.tv_nsec = (timeout_us % 1000000) * 1000;
+    bthread::g_worker_idle_fn = fn;
+    bthread::g_worker_idle_timeout = wait_time;
+    return 0;
+}
+
 int bthread_set_create_span_func(void* (*func)()) {
     if (func == NULL) {
         return EINVAL;

src/bthread/parking_lot.h

@@ -64,15 +64,15 @@ class BAIDU_CACHELINE_ALIGNMENT ParkingLot {
 
     // Wait for tasks.
     // If the `expected_state' does not match, wait() may finish directly.
-    void wait(const State& expected_state) {
+    void wait(const State& expected_state, const timespec* timeout = NULL) {
         if (get_state().val != expected_state.val) {
             // Fast path, no need to futex_wait.
             return;
         }
         if (_no_signal_when_no_waiter) {
             _waiter_num.fetch_add(1, butil::memory_order_relaxed);
         }
-        futex_wait_private(&_pending_signal, expected_state.val, NULL);
+        futex_wait_private(&_pending_signal, expected_state.val, timeout);
         if (_no_signal_when_no_waiter) {
             _waiter_num.fetch_sub(1, butil::memory_order_relaxed);
         }

src/bthread/task_control.cpp

@@ -67,6 +67,10 @@ extern BAIDU_THREAD_LOCAL TaskGroup* tls_task_group;
 void (*g_worker_startfn)() = NULL;
 void (*g_tagged_worker_startfn)(bthread_tag_t) = NULL;
 
+// Worker idle callback configuration
+bool (*g_worker_idle_fn)(void) = NULL;
+timespec g_worker_idle_timeout = {0, 0};
+
 // May be called in other modules to run startfn in non-worker pthreads.
 void run_worker_startfn() {
     if (g_worker_startfn) {
@@ -80,6 +84,16 @@ void run_tagged_worker_startfn(bthread_tag_t tag) {
     }
 }
 
+// Run the idle callback if registered.
+// Returns true if callback returned true (work was done), false otherwise.
+// The callback should access per-worker resources (e.g., io_uring) via TLS.
+bool run_worker_idle_fn() {
+    if (g_worker_idle_fn) {
+        return g_worker_idle_fn();
+    }
+    return false;
+}
+
 struct WorkerThreadArgs {
     WorkerThreadArgs(TaskControl* _c, bthread_tag_t _t) : c(_c), tag(_t) {}
     TaskControl* c;

src/bthread/task_group.cpp

@@ -78,6 +78,11 @@ BAIDU_VOLATILE_THREAD_LOCAL(void*, tls_unique_user_ptr, NULL);
 
 const TaskStatistics EMPTY_STAT = { 0, 0, 0 };
 
+// Defined in task_control.cpp
+extern bool run_worker_idle_fn();
+extern bool (*g_worker_idle_fn)(void);
+extern timespec g_worker_idle_timeout;
+
 void* (*g_create_span_func)() = NULL;
 
 void* run_create_span_func() {
@@ -167,7 +172,17 @@ bool TaskGroup::wait_task(bthread_t* tid) {
         if (_last_pl_state.stopped()) {
             return false;
         }
-        _pl->wait(_last_pl_state);
+        // Instead of waiting for signal, we shall wake up if there's a user idle task here.
+        // To avoid the current task never wake and missed the user's idle task.
+        if (g_worker_idle_fn) {
+            // If we successfuly finished the task, we shall not wait, and start next loop.
+            if (run_worker_idle_fn()) {
+                return true;
+            }
+            _pl->wait(_last_pl_state, &g_worker_idle_timeout);
+        } else {
+            _pl->wait(_last_pl_state);
+        }
         if (steal_task(tid)) {
             return true;
         }
@@ -179,7 +194,17 @@ bool TaskGroup::wait_task(bthread_t* tid) {
         if (steal_task(tid)) {
             return true;
         }
-        _pl->wait(st);
+        // Instead of waiting for signal, we shall wake up if there's a user idle task here.
+        // To avoid the current task never wake and missed the user's idle task.
+        if (g_worker_idle_fn) {
+            if (run_worker_idle_fn()) {
+                // If we successfuly finished the task, we shall not wait, and start next loop.
+                return true;
+            }
+            _pl->wait(st, &g_worker_idle_timeout);
+        } else {
+            _pl->wait(st);
+        }
 #endif
     } while (true);
 }

src/bthread/unstable.h

@@ -92,6 +92,17 @@ extern int bthread_set_worker_startfn(void (*start_fn)());
 // Add a startup function with tag
 extern int bthread_set_tagged_worker_startfn(void (*start_fn)(bthread_tag_t));
 
+// Set a callback to run when a worker has no task to run.
+// If the callback returns true, it means some work is done and the worker
+// should check the runqueue again immediately.
+// The callback should access per-worker resources (e.g., io_uring) via TLS.
+// Users can initialize per-worker resources in bthread_set_worker_startfn().
+// |fn|: The callback function. Pass NULL to clear the callback.
+// |timeout_us|: The timeout for waiting if the callback returns false.
+//               0 is not acceptable when setting a callback (but ignored when clearing).
+// Returns 0 on success, error code otherwise.
+extern int bthread_set_worker_idle_callback(bool (*fn)(void), uint64_t timeout_us);
+
 // Add a create span function
 extern int bthread_set_create_span_func(void* (*func)());
 

test/bthread_idle_unittest.cpp

@@ -0,0 +1,127 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+#include <bthread/bthread.h>
+#include <bthread/unstable.h>
+#include <butil/logging.h>
+#include <butil/time.h>
+#include <set>
+#include <mutex>
+
+namespace {
+
+// Mock context to simulate per-thread state (e.g., io_uring ring)
+struct MockWorkerContext {
+    int worker_id;
+    int poll_count;
+
+    MockWorkerContext() : worker_id(-1), poll_count(0) {}
+};
+
+// Thread-local storage to simulate "shared-nothing" architecture
+// In a real scenario, this would hold something like an io_uring instance.
+static __thread MockWorkerContext* tls_context = nullptr;
+
+// Set to collect all unique worker IDs we've seen
+static std::set<int> observed_worker_ids;
+static std::mutex stats_mutex;
+
+// The idle callback function
+// Access per-worker resources via TLS (simulating io_uring per worker)
+bool MockIdlePoller() {
+    if (!tls_context) {
+        tls_context = new MockWorkerContext();
+        // Use pthread_self or a counter to assign a unique ID
+        static std::atomic<int> global_worker_counter(0);
+        tls_context->worker_id = global_worker_counter.fetch_add(1);
+        
+        std::lock_guard<std::mutex> lock(stats_mutex);
+        observed_worker_ids.insert(tls_context->worker_id);
+        LOG(INFO) << "Worker thread " << pthread_self() << " initialized with ID " << tls_context->worker_id;
+    }
+
+    tls_context->poll_count++;
+    
+    // Simulate some work occasionally to wake up the worker immediately
+    // For this test, we mostly want to verify it runs and has correct context
+    if (tls_context->poll_count % 100 == 0) {
+        return true; // Pretend we found work
+    }
+
+    return false; // Sleep with timeout
+}
+
+class IdleCallbackTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        // Reset state
+        observed_worker_ids.clear();
+    }
+
+    void TearDown() override {
+        // Clean up global callback to avoid affecting other tests
+        bthread_set_worker_idle_callback(nullptr, 0);
+    }
+};
+
+void* dummy_task(void* arg) {
+    bthread_usleep(1000); // Sleep 1ms to allow workers to go idle
+    return nullptr;
+}
+
+TEST_F(IdleCallbackTest, WorkerIsolationAndExecution) {
+    // 1. Set the idle callback with a short timeout (e.g., 1ms)
+    ASSERT_EQ(0, bthread_set_worker_idle_callback(MockIdlePoller, 1000));
+
+    // 2. Determine number of workers (concurrency)
+    int concurrency = bthread_getconcurrency();
+    LOG(INFO) << "Current concurrency: " << concurrency;
+
+    // 3. Create enough bthreads to ensure all workers are activated at least once
+    // but also give them time to become idle.
+    std::vector<bthread_t> tids;
+    for (int i = 0; i < concurrency * 2; ++i) {
+        bthread_t tid;
+        bthread_start_background(&tid, nullptr, dummy_task, nullptr);
+        tids.push_back(tid);
+    }
+
+    // 4. Wait for all tasks to complete
+    for (bthread_t tid : tids) {
+        bthread_join(tid, nullptr);
+    }
+
+    // 5. Sleep a bit to ensure all workers have had a chance to hit the idle loop
+    usleep(50 * 1000); // 50ms
+
+    // 6. Verify results
+    std::lock_guard<std::mutex> lock(stats_mutex);
+    LOG(INFO) << "Observed " << observed_worker_ids.size() << " unique worker contexts.";
+    
+    // We expect at least one worker to have initialized its context.
+    // In a highly concurrent test environment, usually most workers will initialize.
+    ASSERT_GT(observed_worker_ids.size(), 0);
+    
+    // Check that we saw different IDs if concurrency > 1 (though not strictly guaranteed 
+    // that ALL workers will run if the OS scheduler is quirky, but >1 is highly likely)
+    if (concurrency > 1) {
+         EXPECT_GT(observed_worker_ids.size(), 1);
+    }
+}
+
+} // namespace