Destroy keytable pool and tls key when Server::Start() fails early

src/brpc/server.cpp

@@ -451,6 +451,8 @@ Server::Server(ProfilerLinker)
 Server::~Server() {
     Stop(0);
     Join();
+    // Stop()/Join() are no-ops when a Start() failed before RUNNING.
+    DestroyPreStartAllocations();
     ClearServices();
     FreeSSLContexts();
     delete _session_local_data_pool;
@@ -744,6 +746,8 @@ struct RevertServerStatus {
         if (s != nullptr) {
             s->Stop(0);
             s->Join();
+            // Stop()/Join() are no-ops when the start failed before RUNNING.
+            s->DestroyPreStartAllocations();
         }
     }
 };
@@ -1337,24 +1341,9 @@ int Server::Join() {
         _session_local_data_pool->Reset(nullptr);
     }
 
-    if (_keytable_pool) {
-        // Destroy _keytable_pool to delete keytables inside. This has to be
-        // done here (before leaving Join) because it's legal for users to
-        // delete bthread keys after Join which makes related objects
-        // in KeyTables undeletable anymore and leaked.
-        CHECK_EQ(0, bthread_keytable_pool_destroy(_keytable_pool));
-        // TODO: Can't delete _keytable_pool which may be accessed by
-        // still-running bthreads (created by the server). The memory is
-        // leaked but servers are unlikely to be started/stopped frequently,
-        // the leak is acceptable in most scenarios.
-        _keytable_pool = nullptr;
-    }
-
-    // Delete tls_key as well since we don't need it anymore.
-    if (_tl_options.tls_key != INVALID_BTHREAD_KEY) {
-        CHECK_EQ(0, bthread_key_delete(_tl_options.tls_key));
-        _tl_options.tls_key = INVALID_BTHREAD_KEY;
-    }
+    // Destroy pooled keytables before leaving Join(): users may delete
+    // bthread keys afterwards, making their data undeletable and leaked.
+    DestroyThreadLocalResources();
 
     // Have to join _derivative_thread, which may assume that server is running
     // and services in server are not mutated, otherwise data race happens
@@ -1370,6 +1359,34 @@ int Server::Join() {
     return 0;
 }
 
+void Server::DestroyPreStartAllocations() {
+    if (_status == RUNNING || _status == STOPPING) {
+        // Handled by Join().
+        return;
+    }
+    DestroyThreadLocalResources();
+}
+
+void Server::DestroyThreadLocalResources() {
+    if (_keytable_pool != nullptr) {
+        // Keep the TLS key valid until its pooled data has been destroyed.
+        // Deleting the key first would disable DestroyServerTLS().
+        CHECK_EQ(0, bthread_keytable_pool_destroy(_keytable_pool));
+        // Don't delete the pool struct. Bthreads created with
+        // attr.keytable_pool == _keytable_pool may outlive shutdown or
+        // a failed start and still call return_keytable() on exit. A
+        // destroyed pool makes them delete their KeyTables directly, but the
+        // struct itself must stay valid. The leak is annotated in
+        // StartInternal().
+        _keytable_pool = nullptr;
+    }
+    // Delete tls_key created in StartInternal().
+    if (_tl_options.tls_key != INVALID_BTHREAD_KEY) {
+        CHECK_EQ(0, bthread_key_delete(_tl_options.tls_key));
+        _tl_options.tls_key = INVALID_BTHREAD_KEY;
+    }
+}
+
 int Server::AddServiceInternal(google::protobuf::Service* service,
                                bool is_builtin_service,
                                const ServiceOptions& svc_opt) {

src/brpc/server.h

@@ -652,6 +652,15 @@ friend class BadMethodService;
 friend class ServerPrivateAccessor;
 friend class PrometheusMetricsService;
 friend class Controller;
+friend struct RevertServerStatus;
+
+    // Destroy the keytable pool and tls key allocated in StartInternal()
+    // when the start failed before RUNNING: Stop()/Join() are no-ops then,
+    // and a retry would overwrite _keytable_pool and leak the previous one.
+    void DestroyPreStartAllocations();
+
+    // Destroy pooled keytables before deleting the server's TLS key.
+    void DestroyThreadLocalResources();
 
     int AddServiceInternal(google::protobuf::Service* service,
                            bool is_builtin_service,

test/brpc_server_unittest.cpp

@@ -31,6 +31,7 @@
 #include "butil/macros.h"
 #include "butil/fd_guard.h"
 #include "butil/files/scoped_file.h"
+#include "bthread/countdown_event.h"
 #include "brpc/socket.h"
 #include "butil/object_pool.h"
 #include "brpc/builtin/version_service.h"
@@ -220,6 +221,24 @@ class EvilService : public test::EchoService {
     const google::protobuf::ServiceDescriptor* _sd;
 };
 
+class CountingTLSDataFactory : public brpc::DataFactory {
+public:
+    CountingTLSDataFactory() : created(0), destroyed(0) {}
+
+    void* CreateData() const override {
+        created.fetch_add(1, butil::memory_order_relaxed);
+        return new int(0);
+    }
+
+    void DestroyData(void* data) const override {
+        delete static_cast<int*>(data);
+        destroyed.fetch_add(1, butil::memory_order_relaxed);
+    }
+
+    mutable butil::atomic<size_t> created;
+    mutable butil::atomic<size_t> destroyed;
+};
+
 class ServerTest : public ::testing::Test{
 protected:
     ServerTest() {};
@@ -307,6 +326,99 @@ TEST_F(ServerTest, sanity) {
     ASSERT_EQ(0, server.Join());
 }
 
+TEST_F(ServerTest, failed_start_cleans_reserved_thread_local_data) {
+    CountingTLSDataFactory factory;
+    brpc::ServerOptions options;
+    options.thread_local_data_factory = &factory;
+    options.reserved_thread_local_data = 4;
+    brpc::Server listener;
+    ASSERT_EQ(0, listener.Start("127.0.0.1:0", nullptr));
+    brpc::Server server;
+
+    for (size_t attempt = 1; attempt <= 2; ++attempt) {
+        ASSERT_EQ(-1, server.Start(listener.listen_address(), &options));
+        EXPECT_TRUE(server._keytable_pool == nullptr);
+        EXPECT_EQ(INVALID_BTHREAD_KEY, server._tl_options.tls_key);
+        EXPECT_EQ(attempt * options.reserved_thread_local_data,
+                  factory.created.load(butil::memory_order_relaxed));
+        EXPECT_EQ(factory.created.load(butil::memory_order_relaxed),
+                  factory.destroyed.load(butil::memory_order_relaxed));
+    }
+
+    // A successful retry must create a fresh key and pool.
+    ASSERT_EQ(0, server.Start("127.0.0.1:0", &options));
+    EXPECT_TRUE(server._keytable_pool != nullptr);
+    EXPECT_NE(INVALID_BTHREAD_KEY, server._tl_options.tls_key);
+    EXPECT_EQ(3 * options.reserved_thread_local_data,
+              factory.created.load(butil::memory_order_relaxed));
+    ASSERT_EQ(0, server.Stop(0));
+    ASSERT_EQ(0, server.Join());
+    EXPECT_EQ(INVALID_BTHREAD_KEY, server._tl_options.tls_key);
+    EXPECT_EQ(factory.created.load(butil::memory_order_relaxed),
+              factory.destroyed.load(butil::memory_order_relaxed));
+}
+
+TEST_F(ServerTest, failed_start_destruction_cleans_thread_local_data) {
+    CountingTLSDataFactory factory;
+    brpc::ServerOptions options;
+    options.thread_local_data_factory = &factory;
+    options.reserved_thread_local_data = 4;
+    brpc::Server listener;
+    ASSERT_EQ(0, listener.Start("127.0.0.1:0", nullptr));
+    {
+        brpc::Server server;
+        ASSERT_EQ(-1, server.Start(listener.listen_address(), &options));
+        // No successful Start(), Stop() or Join() before destruction.
+    }
+    EXPECT_EQ(options.reserved_thread_local_data,
+              factory.created.load(butil::memory_order_relaxed));
+    EXPECT_EQ(factory.created.load(butil::memory_order_relaxed),
+              factory.destroyed.load(butil::memory_order_relaxed));
+}
+
+static bthread_key_t g_init_child_key = INVALID_BTHREAD_KEY;
+static bthread_t g_init_child = INVALID_BTHREAD;
+
+static void* InitChildOutlivingFailedStart(void* arg) {
+    // Wait until the failed Start() has returned and the Server is destroyed.
+    EXPECT_EQ(0, static_cast<bthread::CountdownEvent*>(arg)->wait());
+    // Create a KeyTable that will be returned to the pool on bthread exit.
+    EXPECT_EQ(0, bthread_setspecific(g_init_child_key, (void*)1));
+    return nullptr;
+}
+
+static bool InitFnSpawningChild(void* arg) {
+    bthread_attr_t attr;
+    // attr.keytable_pool is the server's _keytable_pool.
+    EXPECT_EQ(0, bthread_getattr(bthread_self(), &attr));
+    return bthread_start_background(&g_init_child, &attr,
+                                    InitChildOutlivingFailedStart,
+                                    arg) == 0;
+}
+
+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;
+    ASSERT_EQ(0, listener.Start("127.0.0.1:0", nullptr));
+    bthread::CountdownEvent server_destroyed;
+    brpc::ServerOptions options;
+    options.bthread_init_fn = InitFnSpawningChild;
+    options.bthread_init_args = &server_destroyed;
+    options.bthread_init_count = 1;
+    {
+        brpc::Server server;
+        // Keep signaling and joining the child even if Start() unexpectedly
+        // succeeds, so it cannot remain blocked on a destroyed event.
+        EXPECT_EQ(-1, server.Start(listener.listen_address(), &options));
+    }
+    server_destroyed.signal();
+    // The child calls return_keytable() on the destroyed (but not deleted)
+    // pool. Deleting the pool on the failure path is a use-after-free here
+    // (reported by ASan).
+    ASSERT_EQ(0, bthread_join(g_init_child, nullptr));
+    ASSERT_EQ(0, bthread_key_delete(g_init_child_key));
+}
+
 TEST_F(ServerTest, invalid_protocol_in_enabled_protocols) {
     butil::EndPoint ep;
     ASSERT_EQ(0, str2endpoint("127.0.0.1:8613", &ep));