fix(server): synchronize bthread initialization handoff

src/brpc/server.cpp

@@ -25,6 +25,7 @@
 #include <google/protobuf/descriptor.h>             // ServiceDescriptor
 #include "idl_options.pb.h"                         // option(idl_support)
 #include "bthread/unstable.h"                       // bthread_keytable_pool_init
+#include "butil/atomicops.h"                        // butil::atomic
 #include "butil/macros.h"                           // ARRAY_SIZE
 #include "butil/fd_guard.h"                         // fd_guard
 #include "butil/logging.h"                          // CHECK
@@ -800,16 +801,16 @@ struct BthreadInitArgs {
     bool (*bthread_init_fn)(void* args); // default: nullptr (do nothing)
     void* bthread_init_args;             // default: nullptr
     bool result;
-    bool done;
-    bool stop;
+    butil::atomic<bool> done;
+    butil::atomic<bool> stop;
     bthread_t th;
 };
 
 static void* BthreadInitEntry(void* void_args) {
     BthreadInitArgs* args = (BthreadInitArgs*)void_args;
     args->result = args->bthread_init_fn(args->bthread_init_args);
-    args->done = true;
-    while (!args->stop) {
+    args->done.store(true, butil::memory_order_release);
+    while (!args->stop.load(butil::memory_order_acquire)) {
         bthread_usleep(1000);
     }
     return nullptr;
@@ -1051,8 +1052,8 @@ int Server::StartInternal(const butil::EndPoint& endpoint,
             init_args[i].bthread_init_fn = _options.bthread_init_fn;
             init_args[i].bthread_init_args = _options.bthread_init_args;
             init_args[i].result = false;
-            init_args[i].done = false;
-            init_args[i].stop = false;
+            init_args[i].done.store(false, butil::memory_order_relaxed);
+            init_args[i].stop.store(false, butil::memory_order_relaxed);
             bthread_attr_t tmp = BTHREAD_ATTR_NORMAL;
             tmp.tag = _options.bthread_tag;
             tmp.keytable_pool = _keytable_pool;
@@ -1063,13 +1064,13 @@ int Server::StartInternal(const butil::EndPoint& endpoint,
         }
         // Wait until all created bthreads finish the init function.
         for (size_t i = 0; i < ncreated; ++i) {
-            while (!init_args[i].done) {
+            while (!init_args[i].done.load(butil::memory_order_acquire)) {
                 bthread_usleep(1000);
             }
         }
         // Stop and join created bthreads.
         for (size_t i = 0; i < ncreated; ++i) {
-            init_args[i].stop = true;
+            init_args[i].stop.store(true, butil::memory_order_release);
         }
         for (size_t i = 0; i < ncreated; ++i) {
             bthread_join(init_args[i].th, nullptr);

test/brpc_server_unittest.cpp

@@ -397,6 +397,55 @@ static bool InitFnSpawningChild(void* arg) {
                                     arg) == 0;
 }
 
+struct BthreadInitHandoffArgs {
+    bthread::CountdownEvent init_entered;
+    bthread::CountdownEvent allow_return;
+    butil::atomic<bool> start_finished;
+    brpc::Server* server;
+    brpc::ServerOptions options;
+    int start_result;
+
+    BthreadInitHandoffArgs()
+        : start_finished(false)
+        , server(nullptr)
+        , start_result(-1) {}
+};
+
+static bool InitFnWaitForTest(void* arg) {
+    BthreadInitHandoffArgs* args = static_cast<BthreadInitHandoffArgs*>(arg);
+    args->init_entered.signal();
+    return args->allow_return.wait() == 0;
+}
+
+static void* StartServerWithBthreadInit(void* arg) {
+    BthreadInitHandoffArgs* args = static_cast<BthreadInitHandoffArgs*>(arg);
+    args->start_result = args->server->Start("127.0.0.1:0", &args->options);
+    args->start_finished.store(true, butil::memory_order_release);
+    return nullptr;
+}
+
+// Under TSAN, the barriers deterministically drive the completion and stop
+// handoff between BthreadInitEntry and Server::StartInternal.
+TEST_F(ServerTest, bthread_init_handoff_is_synchronized) {
+    BthreadInitHandoffArgs args;
+    brpc::Server server;
+    args.server = &server;
+    args.options.bthread_init_fn = InitFnWaitForTest;
+    args.options.bthread_init_args = &args;
+    args.options.bthread_init_count = 1;
+
+    pthread_t tid;
+    ASSERT_EQ(0, pthread_create(&tid, nullptr,
+                                StartServerWithBthreadInit, &args));
+    ASSERT_EQ(0, args.init_entered.wait());
+    EXPECT_FALSE(args.start_finished.load(butil::memory_order_acquire));
+    args.allow_return.signal();
+    ASSERT_EQ(0, pthread_join(tid, nullptr));
+    ASSERT_EQ(0, args.start_result);
+    ASSERT_EQ(0, server.Stop(0));
+    ASSERT_EQ(0, server.Join());
+}
+
 TEST_F(ServerTest, failed_start_keeps_keytable_pool_for_init_bthreads) {
     ASSERT_EQ(0, bthread_key_create(&g_init_child_key, nullptr));
     brpc::Server listener;