From 4e5d91b970ac12fdd2c6ee2c370d78fdff7a4112 Mon Sep 17 00:00:00 2001
From: Guogang Li <guogang@google.com>
Date: Thu, 22 Jun 2023 10:04:33 -0700
Subject: [PATCH] Fixed the crash when handling keep alive packet

PiperOrigin-RevId: 542592058
---
 connections/implementation/endpoint_manager.cc | 16 +++++++++++++++-
 connections/implementation/endpoint_manager.h  |  1 +
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/connections/implementation/endpoint_manager.cc b/connections/implementation/endpoint_manager.cc
index 5ad88e8d..6b3b76ea 100644
--- a/connections/implementation/endpoint_manager.cc
+++ b/connections/implementation/endpoint_manager.cc
@@ -21,6 +21,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/time/time.h"
 #include "connections/implementation/analytics/throughput_recorder.h"
 #include "connections/implementation/endpoint_channel.h"
 #include "connections/implementation/offline_frames.h"
@@ -730,6 +731,15 @@ EndpointManager::EndpointState::~EndpointState() {
     MutexLock lock(keep_alive_waiter_mutex_.get());
     keep_alive_waiter_->Notify();
   }
+
+  if (keep_alive_latch_ != nullptr) {
+    ExceptionOr<bool> result = keep_alive_latch_->Await(absl::Seconds(1));
+    if (result.ok() && result.result()) {
+      NEARBY_LOGS(INFO) << "Keep alive thread exit gracefully.";
+    } else {
+      NEARBY_LOGS(WARNING) << "Keep alive thread did not exit gracefully.";
+    }
+  }
 }
 
 void EndpointManager::EndpointState::StartEndpointReader(Runnable&& runnable) {
@@ -738,11 +748,15 @@ void EndpointManager::EndpointState::StartEndpointReader(Runnable&& runnable) {
 
 void EndpointManager::EndpointState::StartEndpointKeepAliveManager(
     std::function<void(Mutex*, ConditionVariable*)> runnable) {
+  keep_alive_latch_ = std::make_unique<CountDownLatch>(1);
   keep_alive_thread_.Execute(
       "keep-alive",
       [runnable, keep_alive_waiter_mutex = keep_alive_waiter_mutex_.get(),
-       keep_alive_waiter = keep_alive_waiter_.get()]() {
+       keep_alive_waiter = keep_alive_waiter_.get(),
+       keep_alive_latch = keep_alive_latch_.get()]() {
         runnable(keep_alive_waiter_mutex, keep_alive_waiter);
+        NEARBY_LOGS(INFO) << "Keep alive thread completed.";
+        keep_alive_latch->CountDown();
       });
 }
 
diff --git a/connections/implementation/endpoint_manager.h b/connections/implementation/endpoint_manager.h
index 922090e9..258aa437 100644
--- a/connections/implementation/endpoint_manager.h
+++ b/connections/implementation/endpoint_manager.h
@@ -202,6 +202,7 @@ class EndpointManager {
     mutable std::unique_ptr<Mutex> keep_alive_waiter_mutex_;
     std::unique_ptr<ConditionVariable> keep_alive_waiter_;
     SingleThreadExecutor keep_alive_thread_;
+    std::unique_ptr<CountDownLatch> keep_alive_latch_;
   };
 
   // RAII accessor for FrameProcessor