[Nearby Connections] Stop 3 AttemptsToConnect over WebRTC for CancellationFlag

CancellationFlags will be used to prevent crashes during the shutdown of Nearby Connections from pending tasks taking too long during the shutdown period. WebRTC using 3 x 10s retries to connect, which means we are potentially waiting for 30 seconds of retries to execute during shutdown (which is longer than the 10s duration alloted for Core shutdown).

By using CancellationFlags, we can prevent the retries occuring during the Shutdown by short-circuiting an in flight AttemptToConnect, and checking for Cancellation before retries.

PiperOrigin-RevId: 539690403

Package.swift

@@ -534,6 +534,7 @@ let package = Package(
         "internal/network/http_status_code_test.cc",
         "internal/test/google3_only/fake_authentication_manager_test.cc",
         "internal/test/fake_clock_test.cc",
+        "internal/test/fake_webrtc.cc",
         "internal/test/fake_timer_test.cc",
         "internal/test/fake_device_info_test.cc",
         "internal/test/fake_task_runner_test.cc",

connections/implementation/mediums/BUILD

@@ -151,9 +151,11 @@ cc_test(
         ":mediums",
         ":utils",
         "//internal/platform:base",
+        "//internal/platform:comm",
         "//internal/platform:test_util",
         "//internal/platform:types",
         "//internal/platform/implementation/g3",  # build_cleaner: keep
+        "//internal/test",
         "@com_github_protobuf_matchers//protobuf-matchers",
         "@com_google_googletest//:gtest_main",
     ],

connections/implementation/mediums/webrtc.cc

@@ -27,6 +27,7 @@
 #include "connections/implementation/mediums/webrtc_socket.h"
 #include "internal/platform/byte_array.h"
 #include "internal/platform/cancelable_alarm.h"
+#include "internal/platform/cancellation_flag_listener.h"
 #include "internal/platform/future.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/mutex_lock.h"
@@ -47,7 +48,10 @@ constexpr absl::Duration kRestartReceiveMessagesDuration = absl::Seconds(60);
 
 }  // namespace
 
-WebRtc::WebRtc() = default;
+WebRtc::WebRtc() : WebRtc(std::make_unique<WebRtcMedium>()) {}
+
+WebRtc::WebRtc(std::unique_ptr<WebRtcMedium> medium)
+    : medium_(std::move(medium)) {}
 
 WebRtc::~WebRtc() {
   // This ensures that all pending callbacks are run before we reset the medium
@@ -65,10 +69,10 @@ WebRtc::~WebRtc() {
 }
 
 const std::string WebRtc::GetDefaultCountryCode() {
-  return medium_.GetDefaultCountryCode();
+  return medium_->GetDefaultCountryCode();
 }
 
-bool WebRtc::IsAvailable() { return medium_.IsValid(); }
+bool WebRtc::IsAvailable() { return medium_->IsValid(); }
 
 bool WebRtc::IsAcceptingConnections(const std::string& service_id) {
   MutexLock lock(&mutex_);
@@ -107,7 +111,7 @@ bool WebRtc::StartAcceptingConnections(const std::string& service_id,
 
   // Create a new SignalingMessenger so that we can communicate w/ Tachyon.
   info.signaling_messenger =
-      medium_.GetSignalingMessenger(self_peer_id.GetId(), location_hint);
+      medium_->GetSignalingMessenger(self_peer_id.GetId(), location_hint);
   if (!info.signaling_messenger->IsValid()) {
     return false;
   }
@@ -198,14 +202,31 @@ WebRtcSocketWrapper WebRtc::Connect(const std::string& service_id,
                                     const WebrtcPeerId& remote_peer_id,
                                     const LocationHint& location_hint,
                                     CancellationFlag* cancellation_flag) {
-  for (int attempts_count = 0; attempts_count < kConnectAttemptsLimit;
-       attempts_count++) {
+  service_id_to_connect_attempts_count_map_[service_id] = 1;
+  while (service_id_to_connect_attempts_count_map_[service_id] <=
+         kConnectAttemptsLimit) {
+    if (cancellation_flag->Cancelled()) {
+      NEARBY_LOGS(WARNING)
+          << "Attempt #"
+          << service_id_to_connect_attempts_count_map_[service_id]
+          << ": Cannot Connect with WebRtc due to cancel.";
+      return WebRtcSocketWrapper();
+    }
+
+    NEARBY_LOGS(INFO) << "Attempt #"
+                      << service_id_to_connect_attempts_count_map_[service_id]
+                      << ": Beginning connection.";
     auto wrapper_result = AttemptToConnect(service_id, remote_peer_id,
                                            location_hint, cancellation_flag);
     if (wrapper_result.IsValid()) {
       return wrapper_result;
     }
+
+    service_id_to_connect_attempts_count_map_[service_id]++;
   }
+
+  NEARBY_LOGS(WARNING) << "Giving up after " << kConnectAttemptsLimit
+                       << " attempts";
   return WebRtcSocketWrapper();
 }
 
@@ -216,6 +237,19 @@ WebRtcSocketWrapper WebRtc::AttemptToConnect(
   info.self_peer_id = WebrtcPeerId::FromRandom();
   Future<WebRtcSocketWrapper> socket_future = info.socket_future;
 
+  // `listener` will go out of scope at the end of `AttemptToConnect`, and this
+  // is expected. This `listener` is tied to `socket_future` which we block on
+  // within this stack call, and will not go out of scope until the attempt
+  // is complete.
+  CancellationFlagListener listener(
+      cancellation_flag, [this, &service_id, &socket_future]() {
+        NEARBY_LOGS(WARNING)
+            << "Attempt # "
+            << service_id_to_connect_attempts_count_map_[service_id]
+            << " to connect with WebRtc stopped due to cancel.";
+        socket_future.SetException({Exception::kFailed});
+      });
+
   {
     MutexLock lock(&mutex_);
     if (!IsAvailable()) {
@@ -226,11 +260,6 @@ WebRtcSocketWrapper WebRtc::AttemptToConnect(
       return WebRtcSocketWrapper();
     }
 
-    if (cancellation_flag->Cancelled()) {
-      NEARBY_LOGS(INFO) << "Cannot connect with WebRtc due to cancel.";
-      return WebRtcSocketWrapper();
-    }
-
     // Create a new ConnectionFlow for this connection attempt.
     std::unique_ptr<ConnectionFlow> connection_flow =
         CreateConnectionFlow(service_id, remote_peer_id);
@@ -244,8 +273,8 @@ WebRtcSocketWrapper WebRtc::AttemptToConnect(
     }
 
     // Create a new SignalingMessenger so that we can communicate over Tachyon.
-    info.signaling_messenger =
-        medium_.GetSignalingMessenger(info.self_peer_id.GetId(), location_hint);
+    info.signaling_messenger = medium_->GetSignalingMessenger(
+        info.self_peer_id.GetId(), location_hint);
     if (!info.signaling_messenger->IsValid()) {
       NEARBY_LOG(
           INFO,
@@ -718,7 +747,7 @@ std::unique_ptr<ConnectionFlow> WebRtc::CreateConnectionFlow(
             });
           }},
       },
-      medium_);
+      *medium_);
 }
 
 void WebRtc::RemoveConnectionFlow(const WebrtcPeerId& remote_peer_id) {

connections/implementation/mediums/webrtc.h

@@ -19,6 +19,7 @@
 
 #include <cstddef>
 #include <functional>
+#include <map>
 #include <memory>
 #include <string>
 
@@ -87,6 +88,14 @@ class WebRtc {
       const location::nearby::connections::LocationHint& location_hint,
       CancellationFlag* cancellation_flag) ABSL_LOCKS_EXCLUDED(mutex_);
 
+ protected:
+  // Use for unit tests only to inject a WebRtcMedium.
+  explicit WebRtc(std::unique_ptr<WebRtcMedium> medium);
+
+  // Used in unit tests to determine how many calls to `AttemptToConnect`
+  // occured during a call to `Connect`, per service id.
+  std::map<std::string, int> service_id_to_connect_attempts_count_map_;
+
  private:
   static constexpr int kConnectAttemptsLimit = 3;
   static constexpr int kRestartAcceptConnectionsLimit = 3;
@@ -225,7 +234,7 @@ class WebRtc {
 
   Mutex mutex_;
 
-  WebRtcMedium medium_;
+  std::unique_ptr<WebRtcMedium> medium_;
 
   // The single thread we throw the potentially blocking work on to.
   ScheduledExecutor single_thread_executor_;

connections/implementation/mediums/webrtc_test.cc

@@ -14,7 +14,9 @@
 
 #include "connections/implementation/mediums/webrtc.h"
 
+#include <memory>
 #include <string>
+#include <utility>
 
 #include "gmock/gmock.h"
 #include "protobuf-matchers/protocol-buffer-matchers.h"
@@ -23,6 +25,8 @@
 #include "internal/platform/listeners.h"
 #include "internal/platform/medium_environment.h"
 #include "internal/platform/mutex_lock.h"
+#include "internal/platform/webrtc.h"
+#include "internal/test/fake_webrtc.h"
 
 namespace nearby {
 namespace connections {
@@ -42,6 +46,16 @@ constexpr FeatureFlags kTestCases[] = {
     },
 };
 
+class TestWebRtc : public WebRtc {
+ public:
+  explicit TestWebRtc(std::unique_ptr<WebRtcMedium> medium)
+      : WebRtc(std::move(medium)) {}
+
+  int connect_attempts_count(std::string service_id) {
+    return service_id_to_connect_attempts_count_map_[service_id];
+  }
+};
+
 class WebRtcTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using MockAcceptedCallback = testing::MockFunction<void(
@@ -404,6 +418,159 @@ TEST_F(WebRtcTest, ContinueAcceptingConnectionsOnComplete) {
   env_.Stop();
 }
 
+// Tests when a CancellationFlag is cancelled during an attempt to
+// `WebRtc::AttemptToConnect` triggered by `WebRtc::Connect`.
+TEST_F(WebRtcTest, CancelDuringConnect) {
+  env_.Start({.webrtc_enabled = true});
+
+  // Enable cancellation flags.
+  env_.SetFeatureFlags(kTestCases[0]);
+
+  WebRtcSocketWrapper receiver_socket, sender_socket;
+  const WebrtcPeerId self_id("self_id");
+  const std::string service_id("NearbySharing");
+  LocationHint location_hint;
+  Future<bool> connected;
+  ByteArray message("message");
+
+  CancellationFlag receiver_flag;
+  std::unique_ptr<WebRtc> receiver = std::make_unique<TestWebRtc>(
+      std::make_unique<FakeWebRtcMedium>(&receiver_flag));
+
+  CancellationFlag sender_flag;
+  std::unique_ptr<WebRtcMedium> sender_medium =
+      std::make_unique<FakeWebRtcMedium>(&sender_flag);
+  FakeWebRtcMedium* fake_sender_medium =
+      static_cast<FakeWebRtcMedium*>(sender_medium.get());
+  auto sender = std::make_unique<TestWebRtc>(std::move(sender_medium));
+
+  // Calls `CancellationFlag::Cancel` during a call to `GetSignalingMessenger`
+  // to simulate the cancellation occuring during an `AttemptToConnect`.
+  fake_sender_medium->TriggerCancellationDuringGetSignalingMessenger();
+
+  receiver->StartAcceptingConnections(
+      service_id, self_id, location_hint,
+      {[&receiver_socket, connected](const std::string& service_id,
+                                     WebRtcSocketWrapper wrapper) mutable {
+        receiver_socket = wrapper;
+        connected.Set(receiver_socket.IsValid());
+      }});
+
+  sender_socket =
+      sender->Connect(service_id, self_id, location_hint, &sender_flag);
+
+  // Since the flag was cancelled during the initial `AttemptToConnect`, except
+  // only one attempt instead of the usual three, because the cancellation flag
+  // should short-circuit the lengthy connection attempts during shutdown.
+  // Because of the way the iteration happens, the check for is cancelled
+  // happens after the counter has already been incremented, but before the
+  // attempt actually occurs.
+  EXPECT_FALSE(sender_socket.IsValid());
+  EXPECT_EQ(2, sender->connect_attempts_count(service_id));
+
+  env_.Stop();
+}
+
+// Tests when a CancellationFlag is cancelled before `WebRtc::Connect` is
+// called.
+TEST_F(WebRtcTest, CancelBeforeConnect) {
+  env_.Start({.webrtc_enabled = true});
+
+  // Enable cancellation flags.
+  env_.SetFeatureFlags(kTestCases[0]);
+
+  WebRtcSocketWrapper receiver_socket, sender_socket;
+  const WebrtcPeerId self_id("self_id");
+  const std::string service_id("NearbySharing");
+  LocationHint location_hint;
+  Future<bool> connected;
+  ByteArray message("message");
+
+  CancellationFlag receiver_flag;
+  std::unique_ptr<WebRtc> receiver = std::make_unique<TestWebRtc>(
+      std::make_unique<FakeWebRtcMedium>(&receiver_flag));
+
+  CancellationFlag sender_flag(true);
+  auto sender = std::make_unique<TestWebRtc>(
+      std::make_unique<FakeWebRtcMedium>(&sender_flag));
+
+  receiver->StartAcceptingConnections(
+      service_id, self_id, location_hint,
+      {[&receiver_socket, connected](const std::string& service_id,
+                                     WebRtcSocketWrapper wrapper) mutable {
+        receiver_socket = wrapper;
+        connected.Set(receiver_socket.IsValid());
+      }});
+
+  sender_socket =
+      sender->Connect(service_id, self_id, location_hint, &sender_flag);
+
+  // Expect an invalid socket from stopping during the first attempt to connect,
+  // because `Connect` returned immediatley when it checked for cancellation.
+  EXPECT_FALSE(sender_socket.IsValid());
+  EXPECT_EQ(1, sender->connect_attempts_count(service_id));
+
+  env_.Stop();
+}
+
+// Tests when a CancellationFlag is cancelled during an attempt to
+// `WebRtc::AttemptToConnect` triggered by `WebRtc::Connect` when multiple
+// `WebRTC::Connect` calls are in flight for multiple service ids.
+TEST_F(WebRtcTest, CancelDuringConnect_MultipleConnect) {
+  env_.Start({.webrtc_enabled = true});
+
+  // Enable cancellation flags.
+  env_.SetFeatureFlags(kTestCases[0]);
+
+  WebRtcSocketWrapper receiver_socket, sender_socket;
+  const WebrtcPeerId self_id("self_id");
+  const std::string ns_service_id("NearbySharing");
+  const std::string ph_service_id("PhoneHub");
+  LocationHint location_hint;
+  Future<bool> connected;
+  ByteArray message("message xyz");
+
+  CancellationFlag receiver_flag;
+  std::unique_ptr<WebRtc> receiver = std::make_unique<TestWebRtc>(
+      std::make_unique<FakeWebRtcMedium>(&receiver_flag));
+
+  CancellationFlag flag;
+  auto sender_medium = std::make_unique<FakeWebRtcMedium>(&flag);
+  FakeWebRtcMedium* fake_sender_medium = sender_medium.get();
+  auto sender = std::make_unique<TestWebRtc>(std::move(sender_medium));
+
+  receiver->StartAcceptingConnections(
+      ns_service_id, self_id, location_hint,
+      {[&receiver_socket, connected](const std::string& ns_service_id,
+                                     WebRtcSocketWrapper wrapper) mutable {
+        receiver_socket = wrapper;
+        connected.Set(receiver_socket.IsValid());
+      }});
+
+  // Simulate a successful connect for the endpoint of NearbySharing.
+  sender_socket = sender->Connect(ns_service_id, self_id, location_hint, &flag);
+  EXPECT_TRUE(sender_socket.IsValid());
+
+  // Calls `CancellationFlag::Cancel` during a call to `GetSignalingMessenger`
+  // to simulate the cancellation occuring during an `AttemptToConnect` for the
+  // endpoint of Phone Hub.
+  fake_sender_medium->TriggerCancellationDuringGetSignalingMessenger();
+  sender_socket = sender->Connect(ph_service_id, self_id, location_hint, &flag);
+  EXPECT_FALSE(sender_socket.IsValid());
+
+  // Since the flag was cancelled during the initial `AttemptToConnect`, except
+  // only one attempt instead of the usual three, because the cancellation flag
+  // should short-circuit the lengthy connection attempts during shutdown.
+  // Because of the way the iteration happens, the check for is cancelled
+  // happens after the counter has already been incremented, but before the
+  // attempt actually occurs. For the successful connect, expect only one
+  // attempt.
+  EXPECT_EQ(1, sender->connect_attempts_count(ns_service_id));
+  EXPECT_EQ(2, sender->connect_attempts_count(ph_service_id));
+
+  env_.Stop();
+}
+
 }  // namespace
 
 }  // namespace mediums

internal/platform/BUILD

@@ -416,6 +416,7 @@ cc_library(
         "//connections:__subpackages__",
         "//fastpair:__subpackages__",
         "//internal/platform/implementation:__subpackages__",
+        "//internal/test:__subpackages__",
         "//presence:__subpackages__",
     ],
     deps = [

internal/platform/webrtc.h

@@ -25,7 +25,7 @@
 
 namespace nearby {
 
-class WebRtcSignalingMessenger final {
+class WebRtcSignalingMessenger {
  public:
   using OnSignalingMessageCallback =
       api::WebRtcSignalingMessenger::OnSignalingMessageCallback;
@@ -35,35 +35,36 @@ class WebRtcSignalingMessenger final {
   explicit WebRtcSignalingMessenger(
       std::unique_ptr<api::WebRtcSignalingMessenger> messenger)
       : impl_(std::move(messenger)) {}
-  ~WebRtcSignalingMessenger() = default;
+  virtual ~WebRtcSignalingMessenger() = default;
   WebRtcSignalingMessenger(WebRtcSignalingMessenger&&) = default;
   WebRtcSignalingMessenger operator=(WebRtcSignalingMessenger&&) = delete;
 
-  bool SendMessage(absl::string_view peer_id, const ByteArray& message) {
+  virtual bool SendMessage(absl::string_view peer_id,
+                           const ByteArray& message) {
     return impl_->SendMessage(peer_id, message);
   }
 
-  bool StartReceivingMessages(
+  virtual bool StartReceivingMessages(
       OnSignalingMessageCallback on_message_callback,
       OnSignalingCompleteCallback on_complete_callback) {
     return impl_->StartReceivingMessages(std::move(on_message_callback),
                                          std::move(on_complete_callback));
   }
 
-  void StopReceivingMessages() { impl_->StopReceivingMessages(); }
+  virtual void StopReceivingMessages() { impl_->StopReceivingMessages(); }
 
-  bool IsValid() const { return impl_ != nullptr; }
+  virtual bool IsValid() const { return impl_ != nullptr; }
 
  private:
   std::unique_ptr<api::WebRtcSignalingMessenger> impl_;
 };
 
-class WebRtcMedium final {
+class WebRtcMedium {
  public:
   using PeerConnectionCallback = api::WebRtcMedium::PeerConnectionCallback;
 
   WebRtcMedium() : impl_(api::ImplementationPlatform::CreateWebRtcMedium()) {}
-  ~WebRtcMedium() = default;
+  virtual ~WebRtcMedium() = default;
   WebRtcMedium(WebRtcMedium&&) = delete;
   WebRtcMedium& operator=(WebRtcMedium&&) = delete;
 
@@ -81,14 +82,14 @@ class WebRtcMedium final {
   }
 
   // Returns a signaling messenger for sending WebRTC signaling messages.
-  std::unique_ptr<WebRtcSignalingMessenger> GetSignalingMessenger(
+  virtual std::unique_ptr<WebRtcSignalingMessenger> GetSignalingMessenger(
       absl::string_view self_id,
       const location::nearby::connections::LocationHint& location_hint) {
     return std::make_unique<WebRtcSignalingMessenger>(
         impl_->GetSignalingMessenger(self_id, location_hint));
   }
 
-  bool IsValid() const { return impl_ != nullptr; }
+  virtual bool IsValid() const { return impl_ != nullptr; }
 
  private:
   std::unique_ptr<api::WebRtcMedium> impl_;

internal/test/BUILD

@@ -21,6 +21,7 @@ cc_library(
         "fake_single_thread_executor.cc",
         "fake_task_runner.cc",
         "fake_timer.cc",
+        "fake_webrtc.cc",
     ],
     hdrs = [
         "fake_clock.h",
@@ -28,6 +29,7 @@ cc_library(
         "fake_single_thread_executor.h",
         "fake_task_runner.h",
         "fake_timer.h",
+        "fake_webrtc.h",
     ],
     copts = [
         "-Ithird_party",
@@ -35,6 +37,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//internal/base:bluetooth_address",
+        "//internal/platform:comm",
         "//internal/platform:types",
         "//internal/platform/implementation:types",
         "@com_google_absl//absl/base:core_headers",

internal/test/fake_webrtc.cc

@@ -0,0 +1,37 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "internal/test/fake_webrtc.h"
+
+#include <memory>
+
+namespace nearby {
+
+FakeWebRtcMedium::FakeWebRtcMedium(CancellationFlag* flag)
+    : WebRtcMedium(), flag_(flag) {}
+
+FakeWebRtcMedium::~FakeWebRtcMedium() = default;
+
+std::unique_ptr<WebRtcSignalingMessenger>
+FakeWebRtcMedium::GetSignalingMessenger(
+    absl::string_view self_id,
+    const location::nearby::connections::LocationHint& location_hint) {
+  if (cancel_during_get_signaling_messenger_) {
+    flag_->Cancel();
+  }
+
+  return WebRtcMedium::GetSignalingMessenger(self_id, location_hint);
+}
+
+}  // namespace nearby

internal/test/fake_webrtc.h

@@ -0,0 +1,53 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef THIRD_PARTY_NEARBY_INTERNAL_TEST_FAKE_WEBRTC_H_
+#define THIRD_PARTY_NEARBY_INTERNAL_TEST_FAKE_WEBRTC_H_
+
+#include <memory>
+
+#include "internal/platform/webrtc.h"
+
+namespace nearby {
+
+class FakeWebRtcMedium : public WebRtcMedium {
+ public:
+  explicit FakeWebRtcMedium(CancellationFlag* flag);
+  FakeWebRtcMedium(FakeWebRtcMedium&&) = delete;
+  FakeWebRtcMedium& operator=(FakeWebRtcMedium&&) = delete;
+  ~FakeWebRtcMedium() override;
+
+  // WebRtcMedium:
+  bool IsValid() const override { return is_valid_; }
+
+  std::unique_ptr<WebRtcSignalingMessenger> GetSignalingMessenger(
+      absl::string_view self_id,
+      const location::nearby::connections::LocationHint& location_hint)
+      override;
+
+  void TriggerCancellationDuringGetSignalingMessenger() {
+    cancel_during_get_signaling_messenger_ = true;
+  }
+
+  void SetIsValid(bool is_valid) { is_valid_ = is_valid; }
+
+ private:
+  CancellationFlag* flag_ = nullptr;
+  bool is_valid_ = true;
+  bool cancel_during_get_signaling_messenger_ = false;
+};
+
+}  // namespace nearby
+
+#endif  // THIRD_PARTY_NEARBY_INTERNAL_TEST_FAKE_WEBRTC_H_