Fix livelock when all TaskGroup runqueues are full (#1038)

src/bthread/task_group.cpp

@@ -333,9 +333,11 @@ void TaskGroup::task_runner(intptr_t skip_remained) {
             g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
         }
 
-#ifndef NDEBUG
+        // Balance the increment done by sched_to() before entering this fresh
+        // stack. Keep the guard held until all remained callbacks finish.
         --g->_sched_recursive_guard;
-#endif
+        g->requeue_pending_task();
+        g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
     }
 
     do {
@@ -733,12 +735,13 @@ extern void CheckBthreadScheSafety();
 
 void TaskGroup::sched_to(TaskGroup** pg, TaskMeta* next_meta) {
     TaskGroup* g = *pg;
-#ifndef NDEBUG
+    // Maintain the guard in release builds as well: push_rq() uses it to
+    // detect remained callbacks. It is balanced below, or by task_runner()
+    // when this switch enters a fresh stack.
     if ((++g->_sched_recursive_guard) > 1) {
         LOG(FATAL) << "Recursively(" << g->_sched_recursive_guard - 1
                    << ") call sched_to(" << g << ")";
     }
-#endif
     // Save errno so that errno is bthread-specific.
     int saved_errno = errno;
     void* saved_unique_user_ptr = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_unique_user_ptr);
@@ -822,14 +825,14 @@ void TaskGroup::sched_to(TaskGroup** pg, TaskMeta* next_meta) {
         g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
     }
 
+    --g->_sched_recursive_guard;
+    g->requeue_pending_task();
+    g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
+
     // Restore errno
     errno = saved_errno;
     // tls_unique_user_ptr probably changed.
     BAIDU_SET_VOLATILE_THREAD_LOCAL(tls_unique_user_ptr, saved_unique_user_ptr);
-
-#ifndef NDEBUG
-    --g->_sched_recursive_guard;
-#endif
     *pg = g;
 }
 
@@ -848,13 +851,16 @@ void TaskGroup::ready_to_run(TaskMeta* meta, bool nosignal) {
     _control->_task_tracer.set_status(TASK_STATUS_READY, meta);
 #endif // BRPC_BTHREAD_TRACER
     push_rq(meta->tid);
+    // push_rq() may switch context and the current task may be resumed by
+    // another group. Use the current group for all subsequent accounting.
+    TaskGroup* g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
     if (nosignal) {
-        ++_num_nosignal;
+        ++g->_num_nosignal;
     } else {
-        const int additional_signal = _num_nosignal;
-        _num_nosignal = 0;
-        _nsignaled += 1 + additional_signal;
-        _control->signal_task(1 + additional_signal, _tag);
+        const int additional_signal = g->_num_nosignal;
+        g->_num_nosignal = 0;
+        g->_nsignaled += 1 + additional_signal;
+        g->_control->signal_task(1 + additional_signal, g->_tag);
     }
 }
 

src/bthread/task_group.h

@@ -83,13 +83,18 @@ class TaskGroup {
     // Suspend caller and run bthread `next_tid' in TaskGroup *pg.
     // Purpose of this function is to avoid pushing `next_tid' to _rq and
     // then being popped by sched(pg), which is not necessary.
+    // On return, *pg is the group currently running the caller, which may
+    // differ from the input group if the caller was stolen while suspended.
     static void sched_to(TaskGroup** pg, TaskMeta* next_meta);
     static void sched_to(TaskGroup** pg, bthread_t next_tid);
     static void exchange(TaskGroup** pg, TaskMeta* next_meta);
 
     // The callback will be run in the beginning of next-run bthread.
     // Can't be called by current bthread directly because it often needs
     // the target to be suspended already.
+    // The callback runs with _sched_recursive_guard held, either in
+    // sched_to()'s remained loop or at task_runner() entry, and must not
+    // directly or indirectly enter sched_to() again.
     typedef void (*RemainedFn)(void*);
     void set_remained(RemainedFn cb, void* arg) {
         _last_context_remained = cb;
@@ -180,8 +185,11 @@ class TaskGroup {
     // Get the meta associate with the task.
     static TaskMeta* address_meta(bthread_t tid);
 
-    // Push a task into _rq, if _rq is full, retry after some time. This
-    // process make go on indefinitely.
+    // Push a task into _rq. If full, run a queued task to make room when
+    // context switching is allowed. Remained callbacks defer a queued task
+    // until the guard is balanced; pthread-mode tasks wait and retry.
+    // This function may switch context. After it returns, the current task
+    // may belong to a different TaskGroup.
     void push_rq(bthread_t tid);
 
     // Returns size of local run queue.
@@ -299,6 +307,7 @@ friend class TaskControl;
     static void asan_task_runner(intptr_t);
 #endif // BUTIL_USE_ASAN
     static void task_runner(intptr_t skip_remained);
+    void requeue_pending_task();
 
     // Callbacks for set_remained()
     static void _release_last_context(void*);
@@ -362,7 +371,12 @@ friend class TaskControl;
     int _remote_num_nosignal{0};
     int _remote_nsignaled{0};
 
+    // Nonzero during a context switch, including remained-callback execution.
+    // Maintained in all build types because push_rq() reads it.
     int _sched_recursive_guard{0};
+    // Task popped from _rq inside a remained callback. Requeue it after
+    // _sched_recursive_guard is balanced. Only accessed by the owning worker.
+    bthread_t _pending_rq_tid{INVALID_BTHREAD};
     // tag of this taskgroup
     bthread_tag_t _tag{BTHREAD_TAG_DEFAULT};
 

src/bthread/task_group_inl.h

@@ -83,22 +83,138 @@ inline void TaskGroup::sched_to(TaskGroup** pg, bthread_t next_tid) {
 }
 
 inline void TaskGroup::push_rq(bthread_t tid) {
-    while (!_rq.push(tid)) {
+    TaskGroup* g = this;
+    while (!g->_rq.push(tid)) {
         // Created too many bthreads: a promising approach is to insert the
         // task into another TaskGroup, but we don't use it because:
         // * There're already many bthreads to run, inserting the bthread
         //   into other TaskGroup does not help.
         // * Insertions into other TaskGroups perform worse when all workers
         //   are busy at creating bthreads (proved by test_input_messenger in
         //   brpc)
-        flush_nosignal_tasks();
-        LOG_EVERY_SECOND(ERROR) << "_rq is full, capacity=" << _rq.capacity();
-        // TODO(gejun): May cause deadlock when all workers are spinning here.
-        // A better solution is to pop and run existing bthreads, however which
-        // make set_remained()-callbacks do context switches and need extensive
-        // reviews on related code.
-        ::usleep(1000);
+        g->flush_nosignal_tasks();
+        LOG_EVERY_SECOND(ERROR) << "_rq is full, capacity=" << g->_rq.capacity();
+        const auto pop_task = [g](bthread_t* t) -> bool {
+#ifndef BTHREAD_FAIR_WSQ
+            return g->_rq.pop(t);
+#else
+            return g->_rq.steal(t);
+#endif
+        };
+        if (g->_sched_recursive_guard != 0) {
+            // Remained callbacks run with the guard held, either in
+            // sched_to()'s remained loop or at task_runner() entry. Re-entering
+            // sched_to() here would be recursive, while waiting here can
+            // livelock if every worker is running a remained callback. Make
+            // room for |tid| and requeue the popped task after the guard is
+            // balanced.
+            CHECK(g->_pending_rq_tid == INVALID_BTHREAD);
+            bthread_t pending_tid;
+            if (!pop_task(&pending_tid)) {
+                continue;
+            }
+            g->_pending_rq_tid = pending_tid;
+            const bool pushed = g->_rq.push(tid);
+            CHECK(pushed) << "push_rq: failed to push after popping a task";
+            return;
+        }
+        if (g->is_current_pthread_task() && !g->is_current_main_task()) {
+            // A pthread-mode bthread runs on _main_stack and cannot be
+            // scheduled out, fall back to waiting.
+            ::usleep(1000);
+            continue;
+        }
+        bthread_t run_tid = 0;
+        if (!pop_task(&run_tid)) {
+            // The queue was drained by stealers between the failed push and
+            // the pop, retrying the push will succeed.
+            continue;
+        }
+        if (g->is_current_main_task() &&
+            address_meta(run_tid)->attr.stack_type == BTHREAD_STACKTYPE_PTHREAD) {
+            // pthread-mode tasks are run inline on _main_stack by the
+            // dispatch loops (run_main_task/task_runner). Switching to one
+            // from the main task itself causes no stack jump and would
+            // silently drop it, so it cannot be run here. But pushing it
+            // straight back makes the LIFO pop() return this same task
+            // forever, hiding any runnable task sitting deeper in the
+            // queue. Hold pthread-mode tasks aside and keep popping until
+            // a runnable one shows up.
+            // Bounded budget: the array must live on the worker's stack
+            // (no allocation on a resource-exhaustion path), and held
+            // tasks are invisible to stealers, so keeping the list short
+            // limits how long they can be delayed. 16 * sizeof(bthread_t)
+            // == 128 bytes == two cache lines.
+            bthread_t pending[16];
+            int npending = 0;
+            pending[npending++] = run_tid;
+            bool found = false;
+            while (npending < 16 && pop_task(&run_tid)) {
+                if (address_meta(run_tid)->attr.stack_type !=
+                        BTHREAD_STACKTYPE_PTHREAD) {
+                    found = true;
+                    break;
+                }
+                pending[npending++] = run_tid;
+            }
+            // Push the held tasks back. The pops above freed at least
+            // |npending| slots, _rq has a single producer (this worker)
+            // and stealers never push, so these pushes must succeed.
+            // Reverse order restores their relative order in the queue.
+            for (int i = npending - 1; i >= 0; --i) {
+                const bool pushed = g->_rq.push(pending[i]);
+                CHECK(pushed) << "push_rq: failed to put back held task";
+            }
+            if (!found) {
+                // Everything within reach is pthread-mode: wait for
+                // stealers to drain the queue.
+                ::usleep(1000);
+                continue;
+            }
+            // Fall through: run_tid is a runnable task, run it inline
+            // below to free a slot for |tid|.
+        }
+        // Run the popped bthread in place instead of spinning, otherwise
+        // all workers may spin here and no one drains the queues any more.
+        // The pop also frees one slot, which guarantees that the deferred
+        // push of the current task succeeds: _rq has a single producer (the
+        // owning worker) and stealers never push.
+        // |args| lives on the stack that sched_to() is about to freeze and is
+        // read by the remained-callback after the switch, so it must stay in
+        // this scope: if it were declared inside the branch below, its slot
+        // could be reused before the callback has finished reading it.
+        ReadyToRunArgs args = { g->tag(), nullptr, false };
+        if (!g->is_current_main_task()) {
+            // Preserve the task being scheduled out. The push is deferred to
+            // a remained-callback: publishing a still-running task into _rq
+            // lets a stealer jump into it and run the same context
+            // concurrently on another worker. The callback runs on the stack
+            // of |run_tid| after sched_to() has frozen the current context,
+            // same as exchange(). Capture the meta now, current_tid() is
+            // |run_tid| by the time the callback runs.
+            // The callback may re-enter push_rq() while _rq is still full;
+            // the guard branch above makes room without switching and defers
+            // the popped task until the guard is balanced.
+            args.meta = g->_cur_meta;
+            g->set_remained(ready_to_run_in_worker_ignoresignal, &args);
+        }
+        // The main task is never queued; ending_sched() resumes it via
+        // _main_tid when the queue drains. Jumping to a pthread-mode task
+        // from a normal bthread resumes the dispatch loop on _main_stack
+        // which runs it inline, same as ending_sched() does.
+        sched_to(&g, run_tid);
+        // Resume here after the popped bthread (and possibly others) ran,
+        // then retry the push of |tid|.
+    }
+}
+
+inline void TaskGroup::requeue_pending_task() {
+    if (_pending_rq_tid == INVALID_BTHREAD) {
+        return;
     }
+    const bthread_t tid = _pending_rq_tid;
+    _pending_rq_tid = INVALID_BTHREAD;
+    push_rq(tid);
 }
 
 inline void TaskGroup::flush_nosignal_tasks_remote() {

test/CMakeLists.txt

@@ -282,6 +282,9 @@ foreach(BTHREAD_UT ${BTHREAD_UNITTESTS})
         brpc-shared-debug
         ${GPERFTOOLS_LIBRARIES})
     add_test(NAME ${BTHREAD_UT_WE} COMMAND ${BTHREAD_UT_WE})
+    if(BTHREAD_UT_WE STREQUAL "bthread_runqueue_unittest")
+        set_tests_properties(${BTHREAD_UT_WE} PROPERTIES TIMEOUT 90)
+    endif()
 endforeach()
 
 # brpc tests