358303772 endpoint manager clean up to reduce chance of countDownLatch crash

cpp/core/internal/base_pcp_handler_test.cc

@@ -15,6 +15,7 @@
 #include "proto/connections/offline_wire_formats.pb.h"
 #include "platform/base/byte_array.h"
 #include "platform/base/exception.h"
+#include "platform/base/medium_environment.h"
 #include "platform/public/count_down_latch.h"
 #include "platform/public/pipe.h"
 #include "proto/connections_enums.pb.h"
@@ -385,9 +386,11 @@ class BasePcpHandlerTest
       .endpoint_distance_changed_cb =
           mock_discovery_listener_.endpoint_distance_changed_cb.AsStdFunction(),
   };
+  MediumEnvironment& env_ = MediumEnvironment::Instance();
 };
 
 TEST_P(BasePcpHandlerTest, ConstructorDestructorWorks) {
+  env_.Start();
   Mediums m;
   EndpointChannelManager ecm;
   EndpointManager em(&ecm);
@@ -395,9 +398,11 @@ TEST_P(BasePcpHandlerTest, ConstructorDestructorWorks) {
   MockPcpHandler pcp_handler(&m, &em, &ecm, &bwu);
   SUCCEED();
   bwu.Shutdown();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, StartAdvertisingChangesState) {
+  env_.Start();
   ClientProxy client;
   Mediums m;
   EndpointChannelManager ecm;
@@ -406,9 +411,11 @@ TEST_P(BasePcpHandlerTest, StartAdvertisingChangesState) {
   MockPcpHandler pcp_handler(&m, &em, &ecm, &bwu);
   StartAdvertising(&client, &pcp_handler);
   bwu.Shutdown();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, StopAdvertisingChangesState) {
+  env_.Start();
   ClientProxy client;
   Mediums m;
   EndpointChannelManager ecm;
@@ -421,9 +428,11 @@ TEST_P(BasePcpHandlerTest, StopAdvertisingChangesState) {
   pcp_handler.StopAdvertising(&client);
   EXPECT_FALSE(client.IsAdvertising());
   bwu.Shutdown();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, StartDiscoveryChangesState) {
+  env_.Start();
   ClientProxy client;
   Mediums m;
   EndpointChannelManager ecm;
@@ -432,9 +441,11 @@ TEST_P(BasePcpHandlerTest, StartDiscoveryChangesState) {
   MockPcpHandler pcp_handler(&m, &em, &ecm, &bwu);
   StartDiscovery(&client, &pcp_handler);
   bwu.Shutdown();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, StopDiscoveryChangesState) {
+  env_.Start();
   ClientProxy client;
   Mediums m;
   EndpointChannelManager ecm;
@@ -447,9 +458,11 @@ TEST_P(BasePcpHandlerTest, StopDiscoveryChangesState) {
   pcp_handler.StopDiscovery(&client);
   EXPECT_FALSE(client.IsDiscovering());
   bwu.Shutdown();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, RequestConnectionChangesState) {
+  env_.Start();
   std::string endpoint_id{"1234"};
   ClientProxy client;
   Mediums m;
@@ -472,9 +485,11 @@ TEST_P(BasePcpHandlerTest, RequestConnectionChangesState) {
   channel_b->Close();
   bwu.Shutdown();
   pcp_handler.DisconnectFromEndpointManager();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, IoError_RequestConnectionFails) {
+  env_.Start();
   std::string endpoint_id{"1234"};
   ClientProxy client;
   Mediums m;
@@ -499,9 +514,11 @@ TEST_P(BasePcpHandlerTest, IoError_RequestConnectionFails) {
   channel_b->Close();
   bwu.Shutdown();
   pcp_handler.DisconnectFromEndpointManager();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, AcceptConnectionChangesState) {
+  env_.Start();
   std::string endpoint_id{"1234"};
   ClientProxy client;
   Mediums m;
@@ -528,9 +545,11 @@ TEST_P(BasePcpHandlerTest, AcceptConnectionChangesState) {
   channel_b->Close();
   bwu.Shutdown();
   pcp_handler.DisconnectFromEndpointManager();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, RejectConnectionChangesState) {
+  env_.Start();
   std::string endpoint_id{"1234"};
   ClientProxy client;
   Mediums m;
@@ -553,9 +572,11 @@ TEST_P(BasePcpHandlerTest, RejectConnectionChangesState) {
   channel_b->Close();
   bwu.Shutdown();
   pcp_handler.DisconnectFromEndpointManager();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, OnIncomingFrameChangesState) {
+  env_.Start();
   std::string endpoint_id{"1234"};
   ClientProxy client;
   Mediums m;
@@ -588,9 +609,11 @@ TEST_P(BasePcpHandlerTest, OnIncomingFrameChangesState) {
   channel_b->Close();
   bwu.Shutdown();
   pcp_handler.DisconnectFromEndpointManager();
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, DestructorIsCalledOnProtocolEndpoint) {
+  env_.Start();
   std::atomic_int destroyed_flag = 0;
   int mediums_count = 0;
   {
@@ -623,9 +646,11 @@ TEST_P(BasePcpHandlerTest, DestructorIsCalledOnProtocolEndpoint) {
     pcp_handler.DisconnectFromEndpointManager();
   }
   EXPECT_EQ(destroyed_flag.load(), mediums_count);
+  env_.Stop();
 }
 
 TEST_P(BasePcpHandlerTest, MultipleMediumsProduceSingleEndpointLostEvent) {
+  env_.Start();
   BooleanMediumSelector allowed = GetParam();
   if (allowed.Count(true) < 2) {
     // Ignore single-medium test cases, and implicit "all mediums" case.
@@ -670,12 +695,14 @@ TEST_P(BasePcpHandlerTest, MultipleMediumsProduceSingleEndpointLostEvent) {
     pcp_handler.DisconnectFromEndpointManager();
   }
   EXPECT_EQ(destroyed_flag.load(), mediums_count);
+  env_.Stop();
 }
 
 INSTANTIATE_TEST_SUITE_P(ParameterizedBasePcpHandlerTest, BasePcpHandlerTest,
                          ::testing::ValuesIn(kTestCases));
 
 TEST_F(BasePcpHandlerTest, InjectEndpoint) {
+  env_.Start();
   std::string service_id{"service"};
   std::string endpoint_id{"ABCD"};
   ClientProxy client;
@@ -727,6 +754,7 @@ TEST_F(BasePcpHandlerTest, InjectEndpoint) {
           .remote_bluetooth_mac_address = ByteArray(kFakeMacAddress),
       });
   bwu.Shutdown();
+  env_.Stop();
 }
 
 }  // namespace

cpp/core/internal/endpoint_manager.cc

@@ -6,10 +6,10 @@
 #include "core/internal/endpoint_channel.h"
 #include "core/internal/offline_frames.h"
 #include "platform/base/exception.h"
+#include "platform/base/feature_flags.h"
 #include "platform/public/count_down_latch.h"
 #include "platform/public/logging.h"
 #include "proto/connections_enums.pb.h"
-
 namespace location {
 namespace nearby {
 namespace connections {
@@ -287,21 +287,19 @@ void EndpointManager::UnregisterFrameProcessor(V1Frame::FrameType frame_type,
 EndpointManager::FrameProcessor* EndpointManager::GetFrameProcessor(
     V1Frame::FrameType frame_type) {
   EndpointManager::FrameProcessor* processor = nullptr;
-  CountDownLatch latch(1);
-  RunOnEndpointManagerThread([this, frame_type, &processor, &latch]() {
-    auto it = frame_processors_.find(frame_type);
-    if (it != frame_processors_.end()) {
-      processor = it->second;
-    }
-    latch.CountDown();
-  });
-  latch.Await();
+  auto it = frame_processors_.find(frame_type);
+  if (it != frame_processors_.end()) {
+    processor = it->second;
+  }
   return processor;
 }
 
 void EndpointManager::EnsureWorkersTerminated(const std::string& endpoint_id) {
+  NEARBY_LOG(ERROR, "EnsureWorkersTerminated for endpoint %s",
+             endpoint_id.c_str());
   auto item = endpoints_.find(endpoint_id);
   if (item != endpoints_.end()) {
+    NEARBY_LOGS(INFO) << "EndpointState found for id: " << endpoint_id;
     // If another instance of data and keep-alive handlers is running, it will
     // terminate soon; we should block until it happens.
     EndpointState& endpoint_state = item->second;
@@ -332,13 +330,21 @@ void EndpointManager::RegisterEndpoint(ClientProxy* client,
   RunOnEndpointManagerThread([this, client, channel = channel.release(),
                               &endpoint_id, &info, &options, &listener,
                               &latch]() {
+    if (endpoints_.contains(endpoint_id)) {
+      NEARBY_LOG(WARNING, "Registing duplicate endpoint %s",
+                 endpoint_id.c_str());
+      if (!FeatureFlags::GetInstance()
+               .GetFlags()
+               .endpoint_manager_ensure_workers_terminated_inside_remove) {
+        EnsureWorkersTerminated(endpoint_id);
+      }
+    }
     // Pass ownership of channel to EndpointChannelManager
     NEARBY_LOG(INFO, "Registering endpoint with channel manager: id=%s",
                endpoint_id.c_str());
     channel_manager_->RegisterChannelForEndpoint(
         client, endpoint_id, std::unique_ptr<EndpointChannel>(channel));
 
-    EnsureWorkersTerminated(endpoint_id);
     EndpointState& endpoint_state =
         endpoints_.emplace(endpoint_id, EndpointState()).first->second;
     endpoint_state.client = client;
@@ -392,6 +398,7 @@ void EndpointManager::RegisterEndpoint(ClientProxy* client,
 
 void EndpointManager::UnregisterEndpoint(ClientProxy* client,
                                          const std::string& endpoint_id) {
+  NEARBY_LOG(ERROR, "UnregisterEndpoint for endpoint %s", endpoint_id.c_str());
   CountDownLatch latch(1);
   RunOnEndpointManagerThread([this, client, endpoint_id, &latch]() {
     RemoveEndpoint(client, endpoint_id,
@@ -428,6 +435,7 @@ std::vector<std::string> EndpointManager::SendPayloadChunk(
 // allow synchronous behavior here it will cause a live lock.
 void EndpointManager::DiscardEndpoint(ClientProxy* client,
                                       const std::string& endpoint_id) {
+  NEARBY_LOG(ERROR, "DiscardEndpoint for endpoint %s", endpoint_id.c_str());
   RunOnEndpointManagerThread([this, client, endpoint_id]() {
     RemoveEndpoint(client, endpoint_id,
                    /*notify=*/
@@ -450,6 +458,7 @@ std::vector<std::string> EndpointManager::SendControlMessage(
 void EndpointManager::RemoveEndpoint(ClientProxy* client,
                                      const std::string& endpoint_id,
                                      bool notify) {
+  NEARBY_LOG(ERROR, "RemoveEndpoint for endpoint %s", endpoint_id.c_str());
   // Unregistering from channel_manager_ will also serve to terminate
   // the dedicated handler and KeepAlive threads we started when we registered
   // this endpoint.
@@ -464,6 +473,11 @@ void EndpointManager::RemoveEndpoint(ClientProxy* client,
     client->OnDisconnected(endpoint_id, notify);
     NEARBY_LOG(INFO, "Removed endpoint; id=%s", endpoint_id.c_str());
   }
+  if (FeatureFlags::GetInstance()
+          .GetFlags()
+          .endpoint_manager_ensure_workers_terminated_inside_remove) {
+    EnsureWorkersTerminated(endpoint_id);
+  }
 }
 
 // @EndpointManagerThread

cpp/platform/base/feature_flags.h

@@ -22,6 +22,9 @@ class FeatureFlags {
     // first one.
     bool disallow_out_of_order_bwu_avail_event = true;
     bool enable_async_bandwidth_upgrade = true;
+    // Let endpoint_manager erase deleted endpoint from endpoints_ inside
+    // function RemoveEndpoint.
+    bool endpoint_manager_ensure_workers_terminated_inside_remove = true;
   };
 
   static const FeatureFlags& GetInstance() {