diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp index 5a29ae236b..19d58ac422 100644 --- a/src/brpc/server.cpp +++ b/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) { diff --git a/src/brpc/server.h b/src/brpc/server.h index 3c68641be5..31dc69af2f 100644 --- a/src/brpc/server.h +++ b/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, diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index cc09761cb7..8073208692 100644 --- a/test/brpc_server_unittest.cpp +++ b/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(data); + destroyed.fetch_add(1, butil::memory_order_relaxed); + } + + mutable butil::atomic created; + mutable butil::atomic 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(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));