modify bthread attribute with tag (#2476)

Co-authored-by: Yang Liming <yangliming02@meituan.com>

src/brpc/acceptor.cpp

@@ -76,8 +76,9 @@ int Acceptor::StartAccept(int listened_fd, int idle_timeout_sec,
         return -1;
     }
     if (idle_timeout_sec > 0) {
-        if (bthread_start_background(&_close_idle_tid, NULL,
-                                     CloseIdleConnections, this) != 0) {
+        bthread_attr_t tmp = BTHREAD_ATTR_NORMAL;
+        tmp.tag = _bthread_tag;
+        if (bthread_start_background(&_close_idle_tid, &tmp, CloseIdleConnections, this) != 0) {
             LOG(FATAL) << "Fail to start bthread";
             return -1;
         }

src/brpc/details/http_message.h

@@ -19,7 +19,7 @@
 #ifndef BRPC_HTTP_MESSAGE_H
 #define BRPC_HTTP_MESSAGE_H
 
-#include <memory>
+#include <memory>                      // std::unique_ptr
 #include <string>                      // std::string
 #include "butil/macros.h"
 #include "butil/iobuf.h"               // butil::IOBuf

src/brpc/server.cpp

@@ -906,6 +906,7 @@ int Server::StartInternal(const butil::EndPoint& endpoint,
             init_args[i].done = false;
             init_args[i].stop = false;
             bthread_attr_t tmp = BTHREAD_ATTR_NORMAL;
+            tmp.tag = _options.bthread_tag;
             tmp.keytable_pool = _keytable_pool;
             if (bthread_start_background(
                     &init_args[i].th, &tmp, BthreadInitEntry, &init_args[i]) != 0) {
@@ -1144,7 +1145,9 @@ int Server::StartInternal(const butil::EndPoint& endpoint,
 
     // Launch _derivative_thread.
     CHECK_EQ(INVALID_BTHREAD, _derivative_thread);
-    if (bthread_start_background(&_derivative_thread, NULL,
+    bthread_attr_t tmp = BTHREAD_ATTR_NORMAL;
+    tmp.tag = _options.bthread_tag;
+    if (bthread_start_background(&_derivative_thread, &tmp,
                                  UpdateDerivedVars, this) != 0) {
         LOG(ERROR) << "Fail to create _derivative_thread";
         return -1;

src/bthread/bthread.cpp

@@ -174,26 +174,32 @@ start_from_non_worker(bthread_t* __restrict tid,
     if (NULL == c) {
         return ENOMEM;
     }
-    TaskGroup* g = NULL;
+    auto tag = BTHREAD_TAG_DEFAULT;
+    if (attr != NULL && attr->tag != BTHREAD_TAG_INVALID) {
+        tag = attr->tag;
+    }
     if (attr != NULL && (attr->flags & BTHREAD_NOSIGNAL)) {
         // Remember the TaskGroup to insert NOSIGNAL tasks for 2 reasons:
         // 1. NOSIGNAL is often for creating many bthreads in batch,
         //    inserting into the same TaskGroup maximizes the batch.
         // 2. bthread_flush() needs to know which TaskGroup to flush.
-        g = tls_task_group_nosignal;
+        auto g = tls_task_group_nosignal;
         if (NULL == g) {
-            g = c->choose_one_group(attr->tag);
+            g = c->choose_one_group(tag);
             tls_task_group_nosignal = g;
         }
         return g->start_background<true>(tid, attr, fn, arg);
     }
-    g = c->choose_one_group(attr ? attr->tag : BTHREAD_TAG_DEFAULT);
-    return g->start_background<true>(tid, attr, fn, arg);
+    return c->choose_one_group(tag)->start_background<true>(tid, attr, fn, arg);
 }
 
-// if tag is default or equal to thread local use thread local task group
+// Meet one of the three conditions, can run in thread local
+// attr is nullptr
+// tag equal to thread local
+// tag equal to BTHREAD_TAG_INVALID
 BUTIL_FORCE_INLINE bool can_run_thread_local(const bthread_attr_t* __restrict attr) {
-    return attr == nullptr || attr->tag == bthread::tls_task_group->tag();
+    return attr == nullptr || attr->tag == bthread::tls_task_group->tag() ||
+           attr->tag == BTHREAD_TAG_INVALID;
 }
 
 struct TidTraits {

src/bthread/task_group.cpp

@@ -40,7 +40,7 @@
 namespace bthread {
 
 static const bthread_attr_t BTHREAD_ATTR_TASKGROUP = {
-    BTHREAD_STACKTYPE_UNKNOWN, 0, NULL, BTHREAD_TAG_DEFAULT };
+    BTHREAD_STACKTYPE_UNKNOWN, 0, NULL, BTHREAD_TAG_INVALID };
 
 static bool pass_bool(const char*, bool) { return true; }
 

src/bthread/types.h

@@ -34,6 +34,7 @@ static const bthread_t INVALID_BTHREAD = 0;
 
 // bthread tag default is 0
 typedef int bthread_tag_t;
+static const bthread_tag_t BTHREAD_TAG_INVALID = -1;
 static const bthread_tag_t BTHREAD_TAG_DEFAULT = 0;
 
 struct sockaddr;
@@ -104,7 +105,7 @@ typedef struct bthread_attr_t {
         stack_type = (stacktype_and_flags & 7);
         flags = (stacktype_and_flags & ~(unsigned)7u);
         keytable_pool = NULL;
-        tag = BTHREAD_TAG_DEFAULT;
+        tag = BTHREAD_TAG_INVALID;
     }
     bthread_attr_t operator|(unsigned other_flags) const {
         CHECK(!(other_flags & 7)) << "flags=" << other_flags;
@@ -122,22 +123,22 @@ typedef struct bthread_attr_t {
 // obvious drawback is that you need more worker pthreads when you have a lot
 // of such bthreads.
 static const bthread_attr_t BTHREAD_ATTR_PTHREAD =
-{ BTHREAD_STACKTYPE_PTHREAD, 0, NULL, BTHREAD_TAG_DEFAULT };
+{ BTHREAD_STACKTYPE_PTHREAD, 0, NULL, BTHREAD_TAG_INVALID };
 
 // bthreads created with following attributes will have different size of
 // stacks. Default is BTHREAD_ATTR_NORMAL.
 static const bthread_attr_t BTHREAD_ATTR_SMALL = {BTHREAD_STACKTYPE_SMALL, 0, NULL,
-                                                  BTHREAD_TAG_DEFAULT};
+                                                  BTHREAD_TAG_INVALID};
 static const bthread_attr_t BTHREAD_ATTR_NORMAL = {BTHREAD_STACKTYPE_NORMAL, 0, NULL,
-                                                   BTHREAD_TAG_DEFAULT};
+                                                   BTHREAD_TAG_INVALID};
 static const bthread_attr_t BTHREAD_ATTR_LARGE = {BTHREAD_STACKTYPE_LARGE, 0, NULL,
-                                                  BTHREAD_TAG_DEFAULT};
+                                                  BTHREAD_TAG_INVALID};
 
 // bthreads created with this attribute will print log when it's started,
 // context-switched, finished.
 static const bthread_attr_t BTHREAD_ATTR_DEBUG = {
     BTHREAD_STACKTYPE_NORMAL, BTHREAD_LOG_START_AND_FINISH | BTHREAD_LOG_CONTEXT_SWITCH, NULL,
-    BTHREAD_TAG_DEFAULT};
+    BTHREAD_TAG_INVALID};
 
 static const size_t BTHREAD_EPOLL_THREAD_NUM = 1;
 static const bthread_t BTHREAD_ATOMIC_INIT = 0;