chenBright commented on code in PR #3647: URL: https://github.com/apache/brpc/pull/3647#discussion_r4358034505
########## test/brpc_span_shutdown_unittest.cpp: ########## @@ -0,0 +1,232 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include <atomic> +#include <cstdlib> +#include <cstring> +#include <thread> +#include <unistd.h> +#include <gflags/gflags.h> +#include <gtest/gtest.h> +#include "butil/files/scoped_temp_dir.h" +#include "butil/time.h" +#include "brpc/span.h" + +namespace brpc { +DECLARE_string(rpcz_database_dir); +bool has_span_db(); +} + +namespace { + +const char kDirectoryEnv[] = "BRPC_SPAN_SHUTDOWN_TEST_DIR"; + +void Require(bool condition, const char* message) { + if (!condition) { + (void)write(STDERR_FILENO, message, strlen(message)); + (void)write(STDERR_FILENO, "\n", 1); + std::_Exit(1); + } +} + +template <typename Predicate> +bool WaitUntil(Predicate predicate, int64_t timeout_ms = 10000) { + const int64_t deadline_ms = butil::monotonic_time_ms() + timeout_ms; + while (!predicate() && butil::monotonic_time_ms() < deadline_ms) { + usleep(1000); + } + return predicate(); +} + +void PrepareChild() { + // Also bound direct shutdown calls and thread joins if a regression hangs. + alarm(15); + const char* directory = getenv(kDirectoryEnv); + Require(directory != nullptr, "missing test database directory"); + brpc::FLAGS_rpcz_database_dir = directory; + brpc::ClearTlsParentSpan(); + Require(!brpc::has_span_db(), "child already has a span database"); +} + +std::shared_ptr<brpc::Span> NewSpan() { + const int64_t now = butil::cpuwide_time_us(); + std::shared_ptr<brpc::Span> span = brpc::Span::CreateClientSpan( + "test.Shutdown/Call", butil::gettimeofday_us() - now); + Require(static_cast<bool>(span), "cannot create span"); + span->set_start_send_us(now); + span->set_start_callback_us(now + 1); + return span; +} + +std::shared_ptr<brpc::Span> StartBlockedDump() { + std::shared_ptr<brpc::Span> span = NewSpan(); + // The dumper publishes the DB before Index() reads info(). This lock keeps + // the dump in flight without adding a hook to the production collector. + Require(pthread_spin_lock(&span->_info_spinlock) == 0, + "cannot lock span annotations"); + brpc::Span::Submit(span, butil::cpuwide_time_us()); + Require(WaitUntil([] { return brpc::has_span_db(); }), + "collector did not open the span database"); + return span; +} + +struct BlockerState { + std::atomic<bool> entered{false}; + std::atomic<bool> release{false}; +}; + +class CollectorBlocker : public bvar::Collected { +public: + explicit CollectorBlocker(BlockerState* state) : _state(state) {} + + void dump_and_destroy(size_t) override { + _state->entered.store(true); + Require(WaitUntil([this] { return _state->release.load(); }), + "collector blocker was not released"); + destroy(); + } + + void destroy() override { delete this; } + bvar::CollectorSpeedLimit* speed_limit() override { return nullptr; } + +private: + BlockerState* _state; +}; + +class SpanShutdownDeathTest : public ::testing::Test { +protected: + void SetUp() override { + // Exec each child: forking an existing collector would leave its + // mutexes and queues without their worker threads. + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + if (getenv(kDirectoryEnv) == nullptr) { + ASSERT_TRUE(_directory.CreateUniqueTempDir()); + ASSERT_EQ(0, setenv(kDirectoryEnv, _directory.path().value().c_str(), 1)); + _owns_directory = true; + } + } + + void TearDown() override { + if (_owns_directory) { + EXPECT_EQ(0, unsetenv(kDirectoryEnv)); + } + } + +private: + // The parent removes the child's files after EXPECT_EXIT has reaped it. + butil::ScopedTempDir _directory; + bool _owns_directory = false; +}; + +TEST_F(SpanShutdownDeathTest, stop_waits_for_inflight_dump) { + EXPECT_EXIT({ + PrepareChild(); + std::shared_ptr<brpc::Span> span = StartBlockedDump(); + std::atomic<bool> stopping{false}; + std::atomic<bool> stopped{false}; + std::thread stopper([&] { + stopping.store(true); + brpc::StopSpanIndexing(); + stopped.store(true); + }); + Require(WaitUntil([&] { return stopping.load(); }), + "shutdown thread did not start"); + Require(!WaitUntil([&] { return stopped.load(); }, 100), + "shutdown returned while a span dump was blocked"); + Require(pthread_spin_unlock(&span->_info_spinlock) == 0, + "cannot release span annotations"); Review Comment: Same as above. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
