[Multiplex] Create ShutdownAll() method to shutdown all the virtual sockets and streams.

PiperOrigin-RevId: 660875041

connections/implementation/mediums/multiplex/multiplex_output_stream.cc

@@ -116,7 +116,7 @@ bool MultiplexOutputStream::Close(const std::string& service_id) {
     return false;
   }
 
-  auto service_id_hash_salt = item->second->GetServiceIdHashSalt();
+  // auto service_id_hash_salt = item->second->GetServiceIdHashSalt();
   item->second->Close();
   if (is_enabled_.Get()) {
     Future<bool> future;
@@ -134,6 +134,24 @@ bool MultiplexOutputStream::Close(const std::string& service_id) {
   return true;
 }
 
+void MultiplexOutputStream::CloseAll() {
+  for (auto& [service_id, virtual_output_stream] : virtual_output_streams_) {
+    if (is_enabled_.Get()) {
+      Future<bool> future;
+      multiplex_writer_.EnqueueToSend(
+          &future,
+          ForDisconnection(service_id,
+                           virtual_output_stream->GetServiceIdHashSalt()),
+          "MultiplexFrame::DISCONNECTION");
+      WaitForResult("MultiplexFrame::DISCONNECTION", &future);
+    }
+    virtual_output_stream->Close();
+  }
+  virtual_output_streams_.clear();
+  physical_writer_->Close();
+  multiplex_writer_.Close();
+}
+
 OutputStream*
 MultiplexOutputStream::CreateVirtualOutputStreamForFirstVirtualSocket(
     const std::string& service_id, const std::string& service_id_hash_salt) {

connections/implementation/mediums/multiplex/multiplex_output_stream.h

@@ -81,6 +81,9 @@ class MultiplexOutputStream {
   // Closes the virtual output stream.
   bool Close(const std::string& service_id);
 
+  // Closes all virtual output streams.
+  void CloseAll();
+
   // Waits for the result of the future.
   Exception WaitForResult(const std::string& method_name, Future<bool>* future);
 

connections/implementation/mediums/multiplex/multiplex_socket.cc

@@ -28,7 +28,6 @@
 #include "connections/implementation/mediums/multiplex/multiplex_output_stream.h"
 #include "connections/implementation/mediums/utils.h"
 #include "internal/platform/base64_utils.h"
-#include "internal/platform/bluetooth_classic.h"
 #include "internal/platform/byte_array.h"
 #include "internal/platform/count_down_latch.h"
 #include "internal/platform/exception.h"
@@ -746,6 +745,40 @@ void MultiplexSocket::Shutdown() {
   NEARBY_LOGS(INFO) << __func__ << " end";
 }
 
+void MultiplexSocket::ShutdownAll() {
+  NEARBY_LOGS(INFO) << __func__ << " start";
+  if (is_shutdown_) {
+    NEARBY_LOGS(WARNING) << __func__ << " Already shutdown";
+    return;
+  }
+
+  CountDownLatch latch(1);
+  RunOffloadThread("VirtualSocketClosed", [this, &latch]() {
+    {
+      MutexLock lock(&virtual_socket_mutex_);
+      multiplex_output_stream_.CloseAll();
+      virtual_sockets_.clear();
+
+      Shutdown();
+    }
+    latch.CountDown();
+  });
+
+  if (!latch.Await(FeatureFlags::GetInstance()
+                      .GetFlags()
+                      .mediums_frame_write_timeout_millis).result() + 200) {
+    NEARBY_LOGS(ERROR) << "Timeout to close virtual socket";
+  }
+
+  NEARBY_LOGS(INFO)
+      << "Shutdown single_thread_offloader_ and physical_reader_thread_";
+  single_thread_offloader_.Shutdown();
+  physical_reader_thread_.Shutdown();
+  NEARBY_LOGS(INFO) << __func__ << " end";
+}
+
+
+
 }  // namespace multiplex
 }  // namespace mediums
 }  // namespace connections

connections/implementation/mediums/multiplex/multiplex_socket.h

@@ -107,10 +107,11 @@ class MultiplexSocket {
   void Shutdown();
   bool IsShutdown() { return is_shutdown_; }
   void SetShutdown(bool is_shutdown) { is_shutdown_ = is_shutdown; }
+  void ShutdownAll();
 
  private:
   explicit MultiplexSocket(std::shared_ptr<MediumSocket> physical_socket);
-  ~MultiplexSocket() = default;
+  ~MultiplexSocket() { ShutdownAll(); };
 
   // Creates the first virtual socket for the service id. The first virtual
   // socket is created by the sender.

connections/implementation/mediums/multiplex/multiplex_socket_test.cc

@@ -373,9 +373,7 @@ TEST(MultiplexSocketTest,
 
   fake_socket_ptr->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();
+  multiplex_socket->ShutdownAll();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0);
 }