Fix the flaky test of multiplex_socket_test.

PiperOrigin-RevId: 657728480

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<CountDownLatch>(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();

connections/implementation/mediums/multiplex/multiplex_socket.h

@@ -19,6 +19,7 @@
 #include <string>
 #include <utility>
 
+#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<CountDownLatch> reader_thread_shutdown_barrier_;
 };
 
 }  // namespace multiplex

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

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<absl::AnyInvocable<void()>> socket_closed_listener) {