epoll bthread deal first

src/brpc/event_dispatcher.cpp

@@ -62,6 +62,7 @@ void InitializeGlobalDispatchers() {
                 FLAGS_usercode_in_pthread ? BTHREAD_ATTR_PTHREAD : BTHREAD_ATTR_NORMAL;
             attr.tag = (BTHREAD_TAG_DEFAULT + i) % FLAGS_task_group_ntags;
             CHECK_EQ(0, g_edisp[i * FLAGS_event_dispatcher_num + j].Start(&attr));
+            bthread_epoll_tid_set(i, g_edisp[i].Tid());
         }
     }
     // This atexit is will be run before g_task_control.stop() because above

src/brpc/event_dispatcher.h

@@ -133,6 +133,8 @@ template <typename T> friend class IOEvent;
     // Returns 0 on success, -1 otherwise and errno is set
     int UnregisterEvent(IOEventDataId event_data_id, int fd, bool pollin);
 
+    bthread_t Tid() const { return _tid; }
+
 private:
     DISALLOW_COPY_AND_ASSIGN(EventDispatcher);
 

src/bthread/bthread.cpp

@@ -625,4 +625,10 @@ uint64_t bthread_cpu_clock_ns(void) {
     return 0;
 }
 
+void bthread_epoll_tid_set(bthread_tag_t tag, bthread_t tid) {
+    CHECK(tag >= BTHREAD_TAG_DEFAULT && tag < FLAGS_task_group_ntags);
+    auto c = bthread::get_task_control();
+    return c->set_group_epoll_tid(tag, tid);
+}
+
 }  // extern "C"

src/bthread/bthread.h

@@ -401,6 +401,9 @@ extern void* bthread_getspecific(bthread_key_t key);
 // Return current bthread tag
 extern bthread_tag_t bthread_self_tag(void);
 
+// set task_groups epoll tid by tag
+extern void bthread_epoll_tid_set(bthread_tag_t tag, bthread_t tid);
+
 // The first call to bthread_once() by any thread in a process, with a given
 // once_control, will call the init_routine() with no arguments. Subsequent
 // calls of bthread_once() with the same once_control will not call the

src/bthread/parking_lot.h

@@ -70,6 +70,7 @@ class BAIDU_CACHELINE_ALIGNMENT ParkingLot {
         _pending_signal.fetch_or(1);
         futex_wake_private(&_pending_signal, 10000);
     }
+
 private:
     // higher 31 bits for signalling, LSB for stopping.
     butil::atomic<int> _pending_signal;

src/bthread/task_control.cpp

@@ -183,6 +183,7 @@ TaskControl::TaskControl()
     , _signal_per_second(&_cumulated_signal_count)
     , _status(print_rq_sizes_in_the_tc, this)
     , _nbthreads("bthread_count")
+    , _priority_qs(FLAGS_task_group_ntags)
     , _pl(FLAGS_task_group_ntags)
 {}
 
@@ -207,6 +208,10 @@ int TaskControl::init(int concurrency) {
         _tagged_worker_usage_second.push_back(new bvar::PerSecond<bvar::PassiveStatus<double>>(
             "bthread_worker_usage", tag_str, _tagged_cumulated_worker_time[i], 1));
         _tagged_nbthreads.push_back(new bvar::Adder<int64_t>("bthread_count", tag_str));
+        if (_priority_qs[i].init(BTHREAD_MAX_CONCURRENCY) != 0) {
+            LOG(FATAL) << "Fail to init _priority_q";
+            return -1;
+        }
     }
 
     // Make sure TimerThread is ready.
@@ -430,6 +435,11 @@ int TaskControl::_destroy_group(TaskGroup* g) {
 
 bool TaskControl::steal_task(bthread_t* tid, size_t* seed, size_t offset) {
     auto tag = tls_task_group->tag();
+
+    if (_priority_qs[tag].steal(tid)) {
+        return true;
+    }
+
     // 1: Acquiring fence is paired with releasing fence in _add_group to
     // avoid accessing uninitialized slot of _groups.
     const size_t ngroup = tag_ngroup(tag).load(butil::memory_order_acquire/*1*/);
@@ -472,6 +482,7 @@ void TaskControl::signal_task(int num_task, bthread_tag_t tag) {
     }
     auto& pl = tag_pl(tag);
     int start_index = butil::fmix64(pthread_numeric_id()) % PARKING_LOT_NUM;
+    // WARNING: This allow some bad case happen when  wait_count is not accurente.
     num_task -= pl[start_index].signal(1);
     if (num_task > 0) {
         for (int i = 1; i < PARKING_LOT_NUM && num_task > 0; ++i) {
@@ -575,4 +586,11 @@ bvar::LatencyRecorder* TaskControl::create_exposed_pending_time() {
     return pt;
 }
 
+void TaskControl::set_group_epoll_tid(bthread_tag_t tag, bthread_t tid) {
+    auto groups = tag_group(tag);
+    const size_t ngroup = tag_ngroup(tag).load(butil::memory_order_acquire);
+    for (size_t i = 0; i < ngroup; i++) {
+        groups[i]->add_epoll_tid(tid);
+    }
+}
 }  // namespace bthread

src/bthread/task_control.h

@@ -30,6 +30,7 @@
 #include <vector>
 #include <array>
 #include <memory>
+#include <unordered_map>
 #include "butil/atomicops.h"                     // butil::atomic
 #include "bvar/bvar.h"                          // bvar::PassiveStatus
 #include "bthread/task_tracer.h"
@@ -96,6 +97,12 @@ friend bthread_t init_for_pthread_stack_trace();
     void stack_trace(std::ostream& os, bthread_t tid);
     std::string stack_trace(bthread_t tid);
 #endif // BRPC_BTHREAD_TRACER
+    // Only deal once when init epoll bthread.
+    void set_group_epoll_tid(bthread_tag_t tag, bthread_t tid);
+
+    void push_priority_q(bthread_tag_t tag, bthread_t tid) {
+        _priority_qs[tag].push(tid);
+    }
 
 private:
     typedef std::array<TaskGroup*, BTHREAD_MAX_CONCURRENCY> TaggedGroups;
@@ -153,13 +160,13 @@ friend bthread_t init_for_pthread_stack_trace();
     std::vector<bvar::PassiveStatus<double>*> _tagged_cumulated_worker_time;
     std::vector<bvar::PerSecond<bvar::PassiveStatus<double>>*> _tagged_worker_usage_second;
     std::vector<bvar::Adder<int64_t>*> _tagged_nbthreads;
+    std::vector<WorkStealingQueue<bthread_t>> _priority_qs;
 
     std::vector<TaggedParkingLot> _pl;
 
 #ifdef BRPC_BTHREAD_TRACER
     TaskTracer _task_tracer;
 #endif // BRPC_BTHREAD_TRACER
-
 };
 
 inline bvar::LatencyRecorder& TaskControl::exposed_pending_time() {

src/bthread/task_group.cpp

@@ -429,12 +429,18 @@ int TaskGroup::start_foreground(TaskGroup** pg,
     } else {
         // NOSIGNAL affects current task, not the new task.
         RemainedFn fn = NULL;
-        if (g->current_task()->about_to_quit) {
+        if (g->cur_epoll_tid()) {
+            fn = priority_to_run;
+        } else if (g->current_task()->about_to_quit) {
             fn = ready_to_run_in_worker_ignoresignal;
         } else {
             fn = ready_to_run_in_worker;
         }
-        ReadyToRunArgs args = { g->_cur_meta, (bool)(using_attr.flags & BTHREAD_NOSIGNAL) };
+        ReadyToRunArgs args = {
+            g->tag(),
+            g->_cur_meta,
+            (bool)(using_attr.flags & BTHREAD_NOSIGNAL)
+        };
         g->set_remained(fn, &args);
         TaskGroup::sched_to(pg, m->tid);
     }
@@ -559,7 +565,6 @@ void TaskGroup::ending_sched(TaskGroup** pg) {
         // Jump to main task if there's no task to run.
         next_tid = g->_main_tid;
     }
-
     TaskMeta* const cur_meta = g->_cur_meta;
     TaskMeta* next_meta = address_meta(next_tid);
     if (next_meta->stack == NULL) {
@@ -798,6 +803,11 @@ void TaskGroup::ready_to_run_in_worker_ignoresignal(void* args_in) {
     return tls_task_group->push_rq(args->meta->tid);
 }
 
+void TaskGroup::priority_to_run(void* args_in) {
+    ReadyToRunArgs* args = static_cast<ReadyToRunArgs*>(args_in);
+    return tls_task_group->control()->push_priority_q(args->tag, args->meta->tid);
+}
+
 struct SleepArgs {
     uint64_t timeout_us;
     bthread_t tid;
@@ -972,7 +982,7 @@ int TaskGroup::interrupt(bthread_t tid, TaskControl* c, bthread_tag_t tag) {
 
 void TaskGroup::yield(TaskGroup** pg) {
     TaskGroup* g = *pg;
-    ReadyToRunArgs args = {  g->_cur_meta, false };
+    ReadyToRunArgs args = { g->tag(), g->_cur_meta, false };
     g->set_remained(ready_to_run_in_worker, &args);
     sched(pg);
 }

src/bthread/task_group.h

@@ -22,6 +22,7 @@
 #ifndef BTHREAD_TASK_GROUP_H
 #define BTHREAD_TASK_GROUP_H
 
+#include <unordered_set>
 #include "butil/time.h"                             // cpuwide_time_ns
 #include "bthread/task_control.h"
 #include "bthread/task_meta.h"                     // bthread_t, TaskMeta
@@ -199,6 +200,10 @@ class TaskGroup {
         total_ns += butil::cputhread_time_ns() - _last_cpu_clock_ns;
         return total_ns;
     }
+    // Thread Unsafe
+    void add_epoll_tid(bthread_t tid) { _epoll_tids.emplace(tid); }
+
+    bool cur_epoll_tid() { return _epoll_tids.count(current_tid()) > 0; }
 
 private:
 friend class TaskControl;
@@ -218,11 +223,13 @@ friend class TaskControl;
     static void _release_last_context(void*);
     static void _add_sleep_event(void*);
     struct ReadyToRunArgs {
+        bthread_tag_t tag;
         TaskMeta* meta;
         bool nosignal;
     };
     static void ready_to_run_in_worker(void*);
     static void ready_to_run_in_worker_ignoresignal(void*);
+    static void priority_to_run(void*);
 
     // Wait for a task to run.
     // Returns true on success, false is treated as permanent error and the
@@ -278,6 +285,7 @@ friend class TaskControl;
 
     // Worker thread id.
     pid_t _tid;
+    std::unordered_set<bthread_t> _epoll_tids;
 };
 
 }  // namespace bthread

src/bthread/task_group_inl.h

@@ -51,7 +51,7 @@ inline void TaskGroup::exchange(TaskGroup** pg, TaskMeta* next_meta) {
     if (g->is_current_pthread_task()) {
         return g->ready_to_run(next_meta);
     }
-    ReadyToRunArgs args = { g->_cur_meta, false };
+    ReadyToRunArgs args = { g->tag(), g->_cur_meta, false };
     g->set_remained((g->current_task()->about_to_quit
                      ? ready_to_run_in_worker_ignoresignal
                      : ready_to_run_in_worker),

test/bthread_setconcurrency_unittest.cpp

@@ -214,7 +214,6 @@ int concurrency_by_tag(int num) {
 
 TEST(BthreadTest, concurrency_by_tag) {
     ASSERT_EQ(concurrency_by_tag(1), false);
-    auto tag_con = bthread_getconcurrency_by_tag(0);
     auto con = bthread_getconcurrency();
     ASSERT_EQ(concurrency_by_tag(con), true);
     ASSERT_EQ(concurrency_by_tag(con + 1), true);