From 1973a803a8090023e7dd7277a0d64c4d3bd6743d Mon Sep 17 00:00:00 2001 From: hai007 Date: Tue, 30 Jul 2024 14:19:54 -0700 Subject: [PATCH] Fix the flaky test of multiplex_socket_test. PiperOrigin-RevId: 657728480 --- .../mediums/multiplex/multiplex_socket.cc | 45 +++++++++++-------- .../mediums/multiplex/multiplex_socket.h | 3 ++ .../multiplex/multiplex_socket_test.cc | 9 ++++ internal/platform/socket.h | 5 +++ 4 files changed, 43 insertions(+), 19 deletions(-) diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.cc b/connections/implementation/mediums/multiplex/multiplex_socket.cc index 933121d8..b7c63d81 100644 --- a/connections/implementation/mediums/multiplex/multiplex_socket.cc +++ b/connections/implementation/mediums/multiplex/multiplex_socket.cc @@ -51,6 +51,10 @@ namespace { // without getting salt from it yet. The fake salt reminds sender to get the // correct socket from `virtualSockets` without remapping it. constexpr absl::string_view kFakeSalt = "RECEIVER_CONDIMENT"; + +// The max duration to wait for the reader thread to stop. +constexpr absl::Duration kTimeoutForReaderThreadStop = absl::Milliseconds(100); + } // namespace using ::location::nearby::mediums::ConnectionResponseFrame; @@ -79,6 +83,9 @@ MultiplexSocket::MultiplexSocket(MediumSocket* physical_socket) : physical_socket_(physical_socket), multiplex_output_stream_{&physical_socket->GetOutputStream(), enabled_}, physical_reader_(&physical_socket->GetInputStream()) { + if (physical_socket->IsFakeSocket()) { + return; + } NEARBY_LOGS(INFO) << "physical_socket_: " << physical_socket_; switch (physical_socket_->GetMedium()) { case Medium::BLUETOOTH: @@ -348,6 +355,7 @@ void MultiplexSocket::StartReaderThread() { "shutdown."; return; } + reader_thread_shutdown_barrier_ = std::make_unique(1); physical_reader_thread_.Execute([this]() { NEARBY_LOGS(INFO) << __func__ << " Reader thread starts."; while (!is_shutdown_) { @@ -382,6 +390,7 @@ void MultiplexSocket::StartReaderThread() { } } if (fail) { + reader_thread_shutdown_barrier_->CountDown(); return; } @@ -660,7 +669,7 @@ void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) { NEARBY_LOGS(INFO) << "Close the physical socket because all virtual " "sockets disconnected."; Shutdown(); - shutdown = true; + // shutdown = true; } } else { NEARBY_LOGS(INFO) << "Virtual socket(" << service_id @@ -673,6 +682,7 @@ void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) { if (!latch.Await(absl::Milliseconds(1000)).result()) { NEARBY_LOGS(ERROR) << "Timeout to close virtual socket"; } + if (shutdown) { NEARBY_LOGS(INFO) << "Shutdown single_thread_offloader_ and physical_reader_thread_"; @@ -732,26 +742,23 @@ void MultiplexSocket::Shutdown() { NEARBY_LOGS(INFO) << __func__ << " Already shutdown"; return; } - { - // MutexLock lock(&virtual_socket_mutex_); - for (auto& [hash_key, virtual_socket] : virtual_sockets_) { - if (virtual_socket != nullptr) { - virtual_socket->Close(); - } - } - virtual_sockets_.clear(); - } multiplex_output_stream_.Shutdown(); - switch (medium_) { - case Medium::BLUETOOTH: - bluetooth_socket_.Close(); - break; - case Medium::UNKNOWN_MEDIUM: - NEARBY_LOGS(INFO) << __func__ << " Unknown medium"; - break; - default: - break; + + if (!physical_socket_->IsFakeSocket()) { + switch (medium_) { + case Medium::BLUETOOTH: + bluetooth_socket_.Close(); + break; + case Medium::UNKNOWN_MEDIUM: + NEARBY_LOGS(INFO) << __func__ << " Unknown medium"; + break; + default: + break; + } + } + if (reader_thread_shutdown_barrier_) { + reader_thread_shutdown_barrier_->Await(kTimeoutForReaderThreadStop); } GetIncomingConnectionCallbacks().clear(); diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.h b/connections/implementation/mediums/multiplex/multiplex_socket.h index 2dcbdb19..c60db654 100644 --- a/connections/implementation/mediums/multiplex/multiplex_socket.h +++ b/connections/implementation/mediums/multiplex/multiplex_socket.h @@ -19,6 +19,7 @@ #include #include +#include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" #include "absl/functional/any_invocable.h" #include "connections/implementation/mediums/multiplex/multiplex_output_stream.h" @@ -27,6 +28,7 @@ #include "internal/platform/ble.h" #include "internal/platform/bluetooth_classic.h" #include "internal/platform/byte_array.h" +#include "internal/platform/count_down_latch.h" #include "internal/platform/future.h" #include "internal/platform/input_stream.h" #include "internal/platform/logging.h" @@ -213,6 +215,7 @@ class MultiplexSocket { // If the socket is already shutdown and no longer in use. bool is_shutdown_ = false; + std::unique_ptr reader_thread_shutdown_barrier_; }; } // namespace multiplex diff --git a/connections/implementation/mediums/multiplex/multiplex_socket_test.cc b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc index 36b812b4..25301449 100644 --- a/connections/implementation/mediums/multiplex/multiplex_socket_test.cc +++ b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc @@ -88,6 +88,9 @@ class FakeSocket : public MediumSocket { writer_2_ = std::move(pipe_2_.second); } + bool IsFakeSocket() override{ + return true; + } InputStream& GetInputStream() override { return *reader_1_; } OutputStream& GetOutputStream() override { return *writer_2_; } Exception Close() override { @@ -198,6 +201,8 @@ TEST(MultiplexSocketTest, CreateSuccessAndReaderThreadStarted) { absl::SleepFor(absl::Milliseconds(100)); fake_socket.reader_1_->Close(); EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 1); + virtual_socket->Close(); + EXPECT_EQ(multiplex_socket_incoming->GetVirtualSocketCount(), 0); } TEST(MultiplexSocketTest, CreateFail_MediumNotSupport) { @@ -368,6 +373,10 @@ TEST(MultiplexSocketTest, fake_socket.reader_1_->Close(); EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 2); + multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2))->Close(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1); + multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_1))->Close(); + EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0); } } // namespace multiplex diff --git a/internal/platform/socket.h b/internal/platform/socket.h index c1a34d64..10736b16 100644 --- a/internal/platform/socket.h +++ b/internal/platform/socket.h @@ -77,6 +77,11 @@ class MediumSocket : public Socket { return false; } + /** Returns true if the socket is a Fake socket for unit test. */ + virtual bool IsFakeSocket() { + return false; + } + /** Adds a listener to be invoked when the socket is closed. */ void AddOnSocketClosedListener( std::unique_ptr> socket_closed_listener) {