Add support for checking all living bthreads

User can check all living bthreads by `curl ip:port/bthreads/all`
or `curl ip:port/bthreads/all?st=1` to show bthread stack trace.
This is an enhancement of the original /bthreads service which
provides a method to check a specified bthread by designated
bthread id, as user has no idea what the bthread id is.

Condisering the performance cost brought by recording the
bthread id on bthread startup and finish, currently this function
is only enabled when BRPC_BTHREAD_TRACER is defined.

Author: ZhengweiZhu

src/brpc/builtin/bthreads_service.cpp

@@ -23,7 +23,10 @@
 #include "brpc/builtin/bthreads_service.h"
 
 namespace bthread {
-extern void print_task(std::ostream& os, bthread_t tid);
+extern void print_task(std::ostream& os, bthread_t tid, bool enable_trace);
+#ifdef BRPC_BTHREAD_TRACER
+extern void print_living_task(std::ostream& os, bool enable_trace);
+#endif // BRPC_BTHREAD_TRACER
 }
 
 
@@ -38,30 +41,42 @@ void BthreadsService::default_method(::google::protobuf::RpcController* cntl_bas
     cntl->http_response().set_content_type("text/plain");
     butil::IOBufBuilder os;
     const std::string& constraint = cntl->http_request().unresolved_path();
-    
     if (constraint.empty()) {
 #ifdef BRPC_BTHREAD_TRACER
-        os << "Use /bthreads/<bthread_id> or /bthreads/<bthread_id>?st=1 for stack trace";
+        // when BRPC_BTHREAD_TRACER enabled, we can trace a specified bthread and
+        // check all living bthread
+        os << "Use /bthreads/<bthread_id> or /bthreads/<bthread_id>?st=1 for stack trace\n";
+        os << "To check all living bthread, use /bthreads/<all> or /bthreads/<all>?st=1 for stack trace\n";
 #else
-        os << "Use /bthreads/<bthread_id>";
+        os << "Use /bthreads/<bthread_id>\n";
 #endif // BRPC_BTHREAD_TRACER
     } else {
-        char* endptr = NULL;
-        bthread_t tid = strtoull(constraint.c_str(), &endptr, 10);
-        if (*endptr == '\0' || *endptr == '/' || *endptr == '?') {
-            ::bthread::print_task(os, tid);
-
+        bool enable_trace = false;
 #ifdef BRPC_BTHREAD_TRACER
             const std::string* st = cntl->http_request().uri().GetQuery("st");
             if (NULL != st && *st == "1") {
-                os << "\nbthread call stack:\n";
-                ::bthread::stack_trace(os, tid);
+                enable_trace = true;
             }
 #endif // BRPC_BTHREAD_TRACER
+        char* endptr = NULL;
+        bthread_t tid = strtoull(constraint.c_str(), &endptr, 10);
+        if (*endptr == '\0' || *endptr == '/' || *endptr == '?') {
+            ::bthread::print_task(os, tid, enable_trace);
+        }
+#ifdef BRPC_BTHREAD_TRACER
+        else if (constraint != "all" && constraint != "all?st=1") {
+            cntl->SetFailed(ENOMETHOD, "path=%s is not a bthread id or all, or all?st=1\n",
+                            constraint.c_str());
         } else {
-            cntl->SetFailed(ENOMETHOD, "path=%s is not a bthread id",
+            ::bthread::print_living_task(os, enable_trace);
+        }
+#else
+        else {
+            cntl->SetFailed(ENOMETHOD, "path=%s is not a bthread id\n",
                             constraint.c_str());
         }
+#endif // BRPC_BTHREAD_TRACER
+
     }
     os.move_to(cntl->response_attachment());
 }

src/brpc/builtin/index_service.cpp

@@ -158,7 +158,11 @@ void IndexService::default_method(::google::protobuf::RpcController* controller,
        << Path("/health", html_addr) << " : Test healthy" << NL
        << Path("/vlog", html_addr) << " : List all VLOG callsites" << NL
        << Path("/sockets", html_addr) << " : Check status of a Socket" << NL
+#ifdef BRPC_BTHREAD_TRACER
+       << Path("/bthreads", html_addr) << " : Check status of a bthread or all living bthreads" << NL
+#else
        << Path("/bthreads", html_addr) << " : Check status of a bthread" << NL
+#endif // BRPC_BTHREAD_TRACER
        << Path("/ids", html_addr) << " : Check status of a bthread_id" << NL
        << Path("/protobufs", html_addr) << " : List all protobuf services and messages" << NL
        << Path("/list", html_addr) << " : json signature of methods" << NL

src/bthread/bthread.cpp

@@ -32,6 +32,7 @@
 #include "bthread/bthread.h"
 
 namespace bthread {
+extern void print_task(std::ostream& os, bthread_t tid, bool enable_trace);
 
 static bool validate_bthread_concurrency(const char*, int32_t val) {
     // bthread_setconcurrency sets the flag on success path which should
@@ -192,6 +193,23 @@ std::string stack_trace(bthread_t tid) {
     }
     return c->stack_trace(tid);
 }
+
+// Print all living (started and not finished) bthread
+void print_living_task(std::ostream& os, bool enable_trace) {
+    TaskControl* c = get_task_control();
+    if (NULL == c) {
+        os << "TaskControl has not been created";
+        return;
+    }
+    std::set<bthread_t> tids = c->get_living_bthreads();
+    if (tids.empty()) {
+        os << "No living bthreads\n";
+        return;
+    }
+    for (auto tid : tids) {
+        print_task(os, tid, enable_trace);
+    }
+}
 #endif // BRPC_BTHREAD_TRACER
 
 static int add_workers_for_each_tag(int num) {

src/bthread/task_control.h

@@ -28,6 +28,7 @@
 #include <signal.h>
 #include <stddef.h>                             // size_t
 #include <vector>
+#include <set>
 #include <array>
 #include <memory>
 #include "butil/atomicops.h"                     // butil::atomic
@@ -40,6 +41,71 @@
 DECLARE_int32(task_group_ntags);
 namespace bthread {
 
+#ifdef BRPC_BTHREAD_TRACER
+class ShardedBthreadSet {
+public:
+    explicit ShardedBthreadSet(size_t shard_count = 16)
+        : _shard_count(shard_count), _shards(shard_count) {}
+
+    void add(bthread_t id) {
+        get_shard(id).add(id);
+    }
+
+    void remove(bthread_t id) {
+        get_shard(id).remove(id);
+    }
+
+    std::set<bthread_t> get_all() {
+        std::set<bthread_t> result;
+        for (auto& shard : _shards) {
+            shard.copy_to(result);
+        }
+        return result;
+    }
+
+private:
+    class Shard {
+    public:
+        Shard() {
+            pthread_spin_init(&_spinlock, PTHREAD_PROCESS_PRIVATE);
+        }
+
+        ~Shard() {
+            pthread_spin_destroy(&_spinlock);
+        }
+
+        Shard(const Shard&) = delete;
+        Shard& operator=(const Shard&) = delete;
+
+        void add(bthread_t id) {
+            BAIDU_SCOPED_LOCK(_spinlock);
+            _bthread_ids.insert(id);
+        }
+
+        void remove(bthread_t id) {
+            BAIDU_SCOPED_LOCK(_spinlock);
+            _bthread_ids.erase(id);
+        }
+
+        void copy_to(std::set<bthread_t>& target) {
+            BAIDU_SCOPED_LOCK(_spinlock);
+            target.insert(_bthread_ids.begin(), _bthread_ids.end());
+        }
+
+    private:
+        pthread_spinlock_t _spinlock;
+        std::set<bthread_t> _bthread_ids;
+    };
+
+    Shard& get_shard(bthread_t id) {
+        return _shards[id % _shard_count];
+    }
+
+    const size_t _shard_count;
+    std::vector<Shard> _shards;
+};
+#endif // BRPC_BTHREAD_TRACER
+
 class TaskGroup;
 
 // Control all task groups
@@ -95,6 +161,9 @@ friend bthread_t init_for_pthread_stack_trace();
     // A stacktrace of bthread can be helpful in debugging.
     void stack_trace(std::ostream& os, bthread_t tid);
     std::string stack_trace(bthread_t tid);
+    std::set<bthread_t> get_living_bthreads() {
+        return _living_bthreads.get_all();
+    }
 #endif // BRPC_BTHREAD_TRACER
 
     void push_priority_queue(bthread_tag_t tag, bthread_t tid) {
@@ -130,6 +199,16 @@ friend bthread_t init_for_pthread_stack_trace();
     bvar::Adder<int64_t>& tag_nworkers(bthread_tag_t tag);
     bvar::Adder<int64_t>& tag_nbthreads(bthread_tag_t tag);
 
+#ifdef BRPC_BTHREAD_TRACER
+    void record_bthread_start(bthread_t tid) {
+        _living_bthreads.add(tid);
+    }
+
+    void record_bthread_finish(bthread_t tid) {
+        _living_bthreads.remove(tid);
+    }
+#endif // BRPC_BTHREAD_TRACER
+
     std::vector<butil::atomic<size_t>> _tagged_ngroup;
     std::vector<TaggedGroups> _tagged_groups;
     butil::Mutex _modify_group_mutex;
@@ -165,6 +244,7 @@ friend bthread_t init_for_pthread_stack_trace();
 
 #ifdef BRPC_BTHREAD_TRACER
     TaskTracer _task_tracer;
+    ShardedBthreadSet _living_bthreads;
 #endif // BRPC_BTHREAD_TRACER
 
 };

src/bthread/task_group.cpp

@@ -36,6 +36,7 @@
 #include "bthread/task_control.h"
 #include "bthread/task_group.h"
 #include "bthread/timer_thread.h"
+#include "bthread/bthread.h"
 
 #ifdef __x86_64__
 #include <x86intrin.h>
@@ -442,6 +443,7 @@ void TaskGroup::task_runner(intptr_t skip_remained) {
             g->_control->_task_tracer.WaitForTracing(m);
         }
         g->_control->_task_tracer.set_status(TASK_STATUS_UNKNOWN, m);
+        g->_control->record_bthread_finish(m->tid);
 #endif // BRPC_BTHREAD_TRACER
 
         g->_control->_nbthreads << -1;
@@ -507,6 +509,7 @@ int TaskGroup::start_foreground(TaskGroup** pg,
     g->_control->tag_nbthreads(g->tag()) << 1;
 #ifdef BRPC_BTHREAD_TRACER
     g->_control->_task_tracer.set_status(TASK_STATUS_CREATED, m);
+    g->_control->record_bthread_start(*th);
 #endif // BRPC_BTHREAD_TRACER
     if (g->is_current_pthread_task()) {
         // never create foreground task in pthread.
@@ -570,6 +573,7 @@ int TaskGroup::start_background(bthread_t* __restrict th,
     _control->tag_nbthreads(tag()) << 1;
 #ifdef BRPC_BTHREAD_TRACER
     _control->_task_tracer.set_status(TASK_STATUS_CREATED, m);
+    _control->record_bthread_start(*th);
 #endif // BRPC_BTHREAD_TRACER
     if (REMOTE) {
         ready_to_run_remote(m, (using_attr.flags & BTHREAD_NOSIGNAL));
@@ -1088,10 +1092,10 @@ void TaskGroup::yield(TaskGroup** pg) {
     sched(pg);
 }
 
-void print_task(std::ostream& os, bthread_t tid) {
+void print_task(std::ostream& os, bthread_t tid, bool enable_trace) {
     TaskMeta* const m = TaskGroup::address_meta(tid);
     if (m == NULL) {
-        os << "bthread=" << tid << " : never existed";
+        os << "bthread=" << tid << " : never existed\n";
         return;
     }
     const uint32_t given_ver = get_version(tid);
@@ -1127,7 +1131,7 @@ void print_task(std::ostream& os, bthread_t tid) {
         }
     }
     if (!matched) {
-        os << "bthread=" << tid << " : not exist now";
+        os << "bthread=" << tid << " : not exist now\n";
     } else {
         os << "bthread=" << tid << " :\nstop=" << stop
            << "\ninterrupted=" << interrupted
@@ -1136,6 +1140,7 @@ void print_task(std::ostream& os, bthread_t tid) {
            << "\narg=" << (void*)arg
            << "\nattr={stack_type=" << attr.stack_type
            << " flags=" << attr.flags
+           << " specified tag=" << attr.tag
            << " keytable_pool=" << attr.keytable_pool
            << "}\nhas_tls=" << has_tls
            << "\nuptime_ns=" << butil::cpuwide_time_ns() - cpuwide_start_ns
@@ -1145,8 +1150,15 @@ void print_task(std::ostream& os, bthread_t tid) {
            << "\nstatus=" << status
            << "\ntraced=" << traced
            << "\nworker_tid=" << worker_tid;
-#else
-           ;
+        if (enable_trace) {
+            os << "\nbthread call stack:\n";
+            stack_trace(os, tid);
+            os << "\n";
+        } else {
+            os << "\n\n";
+        }
+ #else
+           << "\n\n";
            (void)status;(void)traced;(void)worker_tid;
 #endif // BRPC_BTHREAD_TRACER
     }