From 1c8395eedef66ebd0f6b8acfe8aacc0ffed9f486 Mon Sep 17 00:00:00 2001
From: hai007 <hais@google.com>
Date: Tue, 14 May 2024 12:45:43 -0700
Subject: [PATCH] Fix the bug that the multiplex socket is not closed properly.

PiperOrigin-RevId: 633675141
---
 .../mediums/multiplex/multiplex_output_stream.cc     |  4 ++--
 .../mediums/multiplex/multiplex_socket.cc            | 12 ++++++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/connections/implementation/mediums/multiplex/multiplex_output_stream.cc b/connections/implementation/mediums/multiplex/multiplex_output_stream.cc
index 89acacec..80c51732 100644
--- a/connections/implementation/mediums/multiplex/multiplex_output_stream.cc
+++ b/connections/implementation/mediums/multiplex/multiplex_output_stream.cc
@@ -254,7 +254,7 @@ void MultiplexOutputStream::MultiplexWriter::StartWriting() {
     }
     {
       MutexLock lock(&writing_mutex_);
-      if (data_queue_.Empty() && is_writing_) {
+      if (data_queue_.Empty() && is_writing_ && !is_closed_) {
         is_writing_ = false;
         NEARBY_LOGS(INFO) << TAG << "Waiting for data_queue_ has data.";
         Exception wait_succeeded = is_writing_cond_.Wait();
@@ -264,8 +264,8 @@ void MultiplexOutputStream::MultiplexWriter::StartWriting() {
               << ": Failure waiting to wait: " << wait_succeeded.value;
           return;
         }
-        if (is_closed_) break;
       }
+      if (is_closed_) break;
     }
   }
   NEARBY_LOGS(INFO) << TAG << "Writing loop stopped.";
diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.cc b/connections/implementation/mediums/multiplex/multiplex_socket.cc
index 1801900c..6d3a7a2e 100644
--- a/connections/implementation/mediums/multiplex/multiplex_socket.cc
+++ b/connections/implementation/mediums/multiplex/multiplex_socket.cc
@@ -438,8 +438,8 @@ void MultiplexSocket::HandleControlFrame(
     const MultiplexControlFrame& frame) {
   switch (frame.control_frame_type()) {
     case MultiplexControlFrame::CONNECTION_REQUEST:
-      RunOffloadThread("CONNECTION_REQUEST", [this, &salted_service_id_hash,
-                                              &service_id_hash_salt] {
+      RunOffloadThread("CONNECTION_REQUEST", [this, salted_service_id_hash,
+                                              service_id_hash_salt] {
         HandleConnectionRequest(salted_service_id_hash, service_id_hash_salt);
       });
       break;
@@ -451,15 +451,15 @@ void MultiplexSocket::HandleControlFrame(
           << ", ConnectionResponseCode: "
           << frame.connection_response_frame().connection_response_code();
 
-      RunOffloadThread("CONNECTION_RESPONSE", [this, &salted_service_id_hash,
-                                               &service_id_hash_salt,
+      RunOffloadThread("CONNECTION_RESPONSE", [this, salted_service_id_hash,
+                                               service_id_hash_salt,
                                                frame = frame] {
         HandleConnectionResponse(salted_service_id_hash, service_id_hash_salt,
                                  frame.connection_response_frame());
       });
       break;
     case MultiplexControlFrame::DISCONNECTION:
-      RunOffloadThread("DISCONNECTION", [this, &salted_service_id_hash] {
+      RunOffloadThread("DISCONNECTION", [this, salted_service_id_hash] {
         HandleDisconnection(salted_service_id_hash);
       });
       break;
@@ -642,7 +642,7 @@ void MultiplexSocket::OnPhysicalSocketClosed() {
 }
 
 void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) {
-  RunOffloadThread("VirtualSocketClosed", [this, &service_id]() {
+  RunOffloadThread("VirtualSocketClosed", [this, service_id]() {
     {
       MutexLock lock(&virtual_socket_mutex_);
       MediumSocket* virtual_socket = GetVirtualSocket(service_id);