Internal change

PiperOrigin-RevId: 364353126

cpp/platform/base/BUILD

@@ -102,6 +102,7 @@ cc_library(
     ],
     hdrs = [
         "cancellation_flag.h",
+        "cancellation_flag_listener.h",
     ],
     visibility = [
         "//core/internal:__subpackages__",

cpp/platform/base/cancellation_flag.cc

@@ -28,25 +28,37 @@ CancellationFlag::CancellationFlag(bool cancelled) {
   cancelled_ = cancelled;
 }
 
-void CancellationFlag::Cancel() {
-  absl::MutexLock lock(mutex_.get());
+CancellationFlag::~CancellationFlag() {
+  listeners_.clear();
+}
 
+void CancellationFlag::Cancel() {
   // Return immediately as no-op if feature flag is not enabled.
   if (!FeatureFlags::GetInstance().GetFlags().enable_cancellation_flag) {
     return;
   }
 
-  if (cancelled_) {
-    // Someone already cancelled. Return immediately.
-    return;
+  absl::flat_hash_set<CancelListener *> listeners;
+  {
+    absl::MutexLock lock(mutex_.get());
+    if (cancelled_) {
+      // Someone already cancelled. Return immediately.
+      return;
+    }
+    cancelled_ = true;
+
+    listeners = listeners_;
+  }
+
+  for (const auto *listener : listeners) {
+    (*listener)();
   }
-  cancelled_ = true;
 }
 
 bool CancellationFlag::Cancelled() const {
   absl::MutexLock lock(mutex_.get());
 
-  // Return falsea as no-op if feature flag is not enabled.
+  // Return false as no-op if feature flag is not enabled.
   if (!FeatureFlags::GetInstance().GetFlags().enable_cancellation_flag) {
     return false;
   }
@@ -54,5 +66,17 @@ bool CancellationFlag::Cancelled() const {
   return cancelled_;
 }
 
+void CancellationFlag::RegisterOnCancelListener(CancelListener *listener) {
+  absl::MutexLock lock(mutex_.get());
+
+  listeners_.emplace(listener);
+}
+
+void CancellationFlag::UnregisterOnCancelListener(CancelListener *listener) {
+  absl::MutexLock lock(mutex_.get());
+
+  listeners_.erase(listener);
+}
+
 }  // namespace nearby
 }  // namespace location

cpp/platform/base/cancellation_flag.h

@@ -17,6 +17,7 @@
 
 #include <memory>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/synchronization/mutex.h"
 
 namespace location {
@@ -26,13 +27,16 @@ namespace nearby {
 // cleaned up as soon as possible.
 class CancellationFlag {
  public:
+  // The listener for cancellation.
+  using CancelListener = std::function<void()>;
+
   CancellationFlag();
   explicit CancellationFlag(bool cancelled);
   CancellationFlag(const CancellationFlag &) = delete;
   CancellationFlag &operator=(const CancellationFlag &) = delete;
   CancellationFlag(CancellationFlag &&) = default;
   CancellationFlag &operator=(CancellationFlag &&) = default;
-  virtual ~CancellationFlag() = default;
+  virtual ~CancellationFlag();
 
   // Set the flag as cancelled.
   void Cancel() ABSL_LOCKS_EXCLUDED(mutex_);
@@ -41,8 +45,33 @@ class CancellationFlag {
   bool Cancelled() const ABSL_LOCKS_EXCLUDED(mutex_);
 
  private:
+  friend class CancellationFlagListener;
+  friend class CancellationFlagPeer;
+
+  // The registration inserts the pointer of caller's listener callback into
+  // `listeners_`, a flat hash set which support the pointer type for hashing
+  // function. It conducts that 2 different pointers might point to the same
+  // callback function which is unusal and should avoid. Hence we make it as
+  // private and use `CancellationFlagListener` as a RAII to wrap the function.
+  // The caller should register listener as lambda or std::function
+  // via `CancellationFlagListener`.
+  void RegisterOnCancelListener(CancelListener *listener)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  // The un-registration erases the pointer of caller's listener callback from
+  // `listeners_`. This is paired to RegisterOnCancelListener which is
+  // guaranteed to be called under `CancellationFlagListener`.
+  void UnregisterOnCancelListener(CancelListener *listener)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  int CancelListenersSize() const ABSL_LOCKS_EXCLUDED(mutex_) {
+    absl::MutexLock lock(mutex_.get());
+    return listeners_.size();
+  }
+
   std::unique_ptr<absl::Mutex> mutex_;
   bool cancelled_ ABSL_GUARDED_BY(mutex_) = false;
+  absl::flat_hash_set<CancelListener *> ABSL_GUARDED_BY(mutex_) listeners_;
 };
 
 }  // namespace nearby

cpp/platform/base/cancellation_flag_listener.h

@@ -0,0 +1,43 @@
+// Copyright 2020 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 PLATFORM_BASE_CANCELLATION_FLAG_LISTENER_H_
+#define PLATFORM_BASE_CANCELLATION_FLAG_LISTENER_H_
+
+#include "platform/base/cancellation_flag.h"
+
+namespace location {
+namespace nearby {
+
+// An RAII mechanism to register CancelListener over a life cycle of medium
+// class.
+class CancellationFlagListener {
+ public:
+  CancellationFlagListener(CancellationFlag* flag,
+                           std::function<void()> listener)
+      : flag_(flag), listener_(std::move(listener)) {
+    flag_->RegisterOnCancelListener(&listener_);
+  }
+
+  ~CancellationFlagListener() { flag_->UnregisterOnCancelListener(&listener_); }
+
+ private:
+  CancellationFlag* flag_;
+  std::function<void()> listener_;
+};
+
+}  // namespace nearby
+}  // namespace location
+
+#endif  // PLATFORM_BASE_CANCELLATION_FLAG_LISTENER_H_

cpp/platform/base/cancellation_flag_test.cc

@@ -14,15 +14,40 @@
 
 #include "platform/base/cancellation_flag.h"
 
+#include <memory>
+
+#include "platform/base/cancellation_flag_listener.h"
 #include "platform/base/feature_flags.h"
 #include "platform/base/medium_environment.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace location {
 namespace nearby {
+
+class CancellationFlagPeer {
+ public:
+  explicit CancellationFlagPeer(CancellationFlag* cancellation_flag)
+      : cancellation_flag_(cancellation_flag) {}
+  void RegisterOnCancelListener(CancellationFlag::CancelListener* listener) {
+    cancellation_flag_->RegisterOnCancelListener(listener);
+  }
+  void UnregisterOnCancelListener(CancellationFlag::CancelListener* listener) {
+    cancellation_flag_->UnregisterOnCancelListener(listener);
+  }
+  int CancelListenersSize() const {
+    return cancellation_flag_->CancelListenersSize();
+  }
+
+ private:
+  CancellationFlag* cancellation_flag_;  // Not owned by CancellationFlagPeer.
+};
+
 namespace {
 
 using FeatureFlags = FeatureFlags::Flags;
+using ::testing::MockFunction;
+using ::testing::StrictMock;
 
 constexpr FeatureFlags kTestCases[] = {
     FeatureFlags{
@@ -64,7 +89,13 @@ TEST_P(CancellationFlagTest, InitialValueAsTrue) {
 }
 
 TEST_P(CancellationFlagTest, CanCancel) {
+  StrictMock<MockFunction<void()>> mock_cancel_callback;
+  CancellationFlag::CancelListener cancel_callback =
+      mock_cancel_callback.AsStdFunction();
+  EXPECT_CALL(mock_cancel_callback, Call)
+      .Times(feature_flags_.enable_cancellation_flag ? 1 : 0);
   CancellationFlag flag;
+  CancellationFlagListener cancellation_flag_listener(&flag, cancel_callback);
   flag.Cancel();
 
   // If FeatureFlag is disabled, return as no-op immediately and
@@ -77,9 +108,105 @@ TEST_P(CancellationFlagTest, CanCancel) {
   EXPECT_TRUE(flag.Cancelled());
 }
 
+TEST_P(CancellationFlagTest, ShouldOnlyCancelOnce) {
+  StrictMock<MockFunction<void()>> mock_cancel_callback;
+  CancellationFlag::CancelListener cancel_callback =
+      mock_cancel_callback.AsStdFunction();
+  EXPECT_CALL(mock_cancel_callback, Call)
+      .Times(feature_flags_.enable_cancellation_flag ? 1 : 0);
+
+  CancellationFlag flag;
+  CancellationFlagListener cancellation_flag_listener(&flag, cancel_callback);
+  flag.Cancel();
+  flag.Cancel();
+  flag.Cancel();
+
+  // If FeatureFlag is disabled, return as no-op immediately and
+  // Cancelled is always false.
+  if (!feature_flags_.enable_cancellation_flag) {
+    EXPECT_FALSE(flag.Cancelled());
+    return;
+  }
+
+  EXPECT_TRUE(flag.Cancelled());
+}
+
+TEST_P(CancellationFlagTest, CannotCancelAfterUnregister) {
+  StrictMock<MockFunction<void()>> mock_cancel_callback;
+  CancellationFlag::CancelListener cancel_callback =
+      mock_cancel_callback.AsStdFunction();
+  EXPECT_CALL(mock_cancel_callback, Call).Times(0);
+
+  CancellationFlag flag;
+  auto cancellation_flag_listener =
+      std::make_unique<CancellationFlagListener>(&flag, cancel_callback);
+  // Release immediately.
+  cancellation_flag_listener.reset();
+  flag.Cancel();
+}
+
 INSTANTIATE_TEST_SUITE_P(ParametrisedCancellationFlagTest, CancellationFlagTest,
                          ::testing::ValuesIn(kTestCases));
 
 }  // namespace
+
+TEST(CancellationFlagTest,
+     CancelMultiplesIfMultiplePointersToTheSameFunctionRegistered) {
+  location::nearby::FeatureFlags::Flags feature_flags_ =
+      location::nearby::FeatureFlags::Flags{
+          .enable_cancellation_flag = true,
+      };
+  MediumEnvironment::Instance().SetFeatureFlags(feature_flags_);
+
+  StrictMock<MockFunction<void()>> mock_cancel_callback;
+  CancellationFlag::CancelListener cancel_callback =
+      mock_cancel_callback.AsStdFunction();
+
+  CancellationFlag::CancelListener *callback_pointer_1 = &cancel_callback;
+  auto callback_pointer_2 =
+      std::make_unique<CancellationFlag::CancelListener>();
+  *callback_pointer_2 = cancel_callback;
+
+  EXPECT_NE(callback_pointer_1, callback_pointer_2.get());
+  EXPECT_CALL(mock_cancel_callback, Call).Times(2);
+
+  CancellationFlag flag;
+  CancellationFlagPeer flag_peer(&flag);
+  flag_peer.RegisterOnCancelListener(callback_pointer_1);
+  flag_peer.RegisterOnCancelListener(callback_pointer_2.get());
+
+  flag.Cancel();
+
+  flag_peer.UnregisterOnCancelListener(callback_pointer_2.get());
+  EXPECT_EQ(1, flag_peer.CancelListenersSize());
+  flag_peer.UnregisterOnCancelListener(callback_pointer_1);
+  EXPECT_EQ(0, flag_peer.CancelListenersSize());
+}
+
+TEST(CancellationFlagTest, RegisteredMultuipleTimesOnlyCancelOnce) {
+  location::nearby::FeatureFlags::Flags feature_flags_ =
+      location::nearby::FeatureFlags::Flags{
+          .enable_cancellation_flag = true,
+      };
+  MediumEnvironment::Instance().SetFeatureFlags(feature_flags_);
+
+  StrictMock<MockFunction<void()>> mock_cancel_callback;
+  CancellationFlag::CancelListener cancel_callback =
+      mock_cancel_callback.AsStdFunction();
+  EXPECT_CALL(mock_cancel_callback, Call).Times(1);
+
+  CancellationFlag flag;
+  CancellationFlagPeer flag_peer(&flag);
+  flag_peer.RegisterOnCancelListener(&cancel_callback);
+  flag_peer.RegisterOnCancelListener(&cancel_callback);
+  EXPECT_EQ(1, flag_peer.CancelListenersSize());
+
+  flag.Cancel();
+
+  flag_peer.UnregisterOnCancelListener(&cancel_callback);
+  EXPECT_EQ(0, flag_peer.CancelListenersSize());
+  flag_peer.UnregisterOnCancelListener(&cancel_callback);
+}
+
 }  // namespace nearby
 }  // namespace location

cpp/platform/impl/g3/ble.cc

@@ -19,6 +19,7 @@
 #include <string>
 
 #include "platform/api/ble.h"
+#include "platform/base/cancellation_flag_listener.h"
 #include "platform/base/logging.h"
 #include "platform/base/medium_environment.h"
 #include "absl/synchronization/mutex.h"
@@ -44,9 +45,7 @@ InputStream& BleSocket::GetInputStream() {
   return remote_socket->GetLocalInputStream();
 }
 
-OutputStream& BleSocket::GetOutputStream() {
-  return GetLocalOutputStream();
-}
+OutputStream& BleSocket::GetOutputStream() { return GetLocalOutputStream(); }
 
 BleSocket* BleSocket::GetRemoteSocket() {
   absl::MutexLock lock(&mutex_);
@@ -354,6 +353,11 @@ std::unique_ptr<api::BleSocket> BleMedium::Connect(
     return {};
   }
 
+  CancellationFlagListener listener(cancellation_flag, [this]() {
+    NEARBY_LOGS(INFO) << "G3 BLE Cancel Connect.";
+    if (server_socket_ != nullptr) server_socket_->Close();
+  });
+
   BlePeripheral peripheral = static_cast<BlePeripheral&>(remote_peripheral);
   auto socket = std::make_unique<BleSocket>(&peripheral);
   // Finally, Request to connect to this socket.

cpp/platform/impl/g3/bluetooth_classic.cc

@@ -18,6 +18,7 @@
 #include <string>
 
 #include "platform/api/bluetooth_classic.h"
+#include "platform/base/cancellation_flag_listener.h"
 #include "platform/base/logging.h"
 #include "platform/base/medium_environment.h"
 #include "platform/impl/g3/bluetooth_adapter.h"
@@ -234,6 +235,11 @@ std::unique_ptr<api::BluetoothSocket> BluetoothClassicMedium::ConnectToService(
     return {};
   }
 
+  CancellationFlagListener listener(cancellation_flag, [&server_socket]() {
+    NEARBY_LOGS(INFO) << "G3 Bluetooth Cancel Connect.";
+    if (server_socket != nullptr) server_socket->Close();
+  });
+
   auto socket = std::make_unique<BluetoothSocket>(&GetAdapter());
   // Finally, Request to connect to this socket.
   if (!server_socket->Connect(*socket)) {

cpp/platform/impl/g3/wifi_lan.cc

@@ -19,6 +19,7 @@
 #include <string>
 
 #include "platform/api/wifi_lan.h"
+#include "platform/base/cancellation_flag_listener.h"
 #include "platform/base/logging.h"
 #include "platform/base/medium_environment.h"
 #include "platform/base/nsd_service_info.h"
@@ -361,6 +362,11 @@ std::unique_ptr<api::WifiLanSocket> WifiLanMedium::Connect(
     return {};
   }
 
+  CancellationFlagListener listener(cancellation_flag, [this]() {
+    NEARBY_LOGS(INFO) << "G3 WifiLan Cancel Connect.";
+    if (server_socket_ != nullptr) server_socket_->Close();
+  });
+
   WifiLanService wifi_lan_service =
       static_cast<WifiLanService&>(remote_wifi_lan_service);
   auto socket = std::make_unique<WifiLanSocket>(&wifi_lan_service);

cpp/platform/public/ble_test.cc

@@ -65,6 +65,7 @@ TEST_P(BleMediumTest, CanStartAcceptingConnectionsAndConnect) {
   std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
   CountDownLatch found_latch(1);
   CountDownLatch accepted_latch(1);
+  CancellationFlag flag;
 
   BlePeripheral* discovered_peripheral = nullptr;
   ble_a.StartScanning(
@@ -102,8 +103,7 @@ TEST_P(BleMediumTest, CanStartAcceptingConnectionsAndConnect) {
   {
     SingleThreadExecutor client_executor;
     client_executor.Execute(
-        [&ble_a, &socket_a, discovered_peripheral, &service_id]() {
-          CancellationFlag flag;
+        [&ble_a, &socket_a, discovered_peripheral, &service_id, &flag]() {
           socket_a = ble_a.Connect(*discovered_peripheral, service_id, &flag);
         });
   }
@@ -127,6 +127,7 @@ TEST_P(BleMediumTest, CanCancelConnect) {
   std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
   CountDownLatch found_latch(1);
   CountDownLatch accepted_latch(1);
+  CancellationFlag flag(true);
 
   BlePeripheral* discovered_peripheral = nullptr;
   ble_a.StartScanning(
@@ -164,8 +165,7 @@ TEST_P(BleMediumTest, CanCancelConnect) {
   {
     SingleThreadExecutor client_executor;
     client_executor.Execute(
-        [&ble_a, &socket_a, discovered_peripheral, &service_id]() {
-          CancellationFlag flag(true);
+        [&ble_a, &socket_a, discovered_peripheral, &service_id, &flag]() {
           socket_a = ble_a.Connect(*discovered_peripheral, service_id, &flag);
         });
   }

cpp/platform/public/bluetooth_classic_test.cc

@@ -108,15 +108,15 @@ TEST_P(BluetoothClassicMediumTest, CanConnectToService) {
   EXPECT_FALSE(socket_a.IsValid());
   EXPECT_FALSE(socket_b.IsValid());
   {
+    CancellationFlag flag;
     SingleThreadExecutor server_executor;
     SingleThreadExecutor client_executor;
-    client_executor.Execute(
-        [this, &socket_a, discovered_device, &service_uuid, &server_socket]() {
-          CancellationFlag flag;
-          socket_a =
-              bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
-          if (!socket_a.IsValid()) server_socket.Close();
-        });
+    client_executor.Execute([this, &socket_a, discovered_device, &service_uuid,
+                             &server_socket, &flag]() {
+      socket_a =
+          bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
+      if (!socket_a.IsValid()) server_socket.Close();
+    });
     server_executor.Execute([&socket_b, &server_socket]() {
       socket_b = server_socket.Accept();
       if (!socket_b.IsValid()) server_socket.Close();
@@ -157,15 +157,15 @@ TEST_P(BluetoothClassicMediumTest, CanCancelConnect) {
   EXPECT_FALSE(socket_a.IsValid());
   EXPECT_FALSE(socket_b.IsValid());
   {
+    CancellationFlag flag(true);
     SingleThreadExecutor server_executor;
     SingleThreadExecutor client_executor;
-    client_executor.Execute(
-        [this, &socket_a, discovered_device, &service_uuid, &server_socket]() {
-          CancellationFlag flag(true);
-          socket_a =
-              bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
-          if (!socket_a.IsValid()) server_socket.Close();
-        });
+    client_executor.Execute([this, &socket_a, discovered_device, &service_uuid,
+                             &server_socket, &flag]() {
+      socket_a =
+          bt_a_->ConnectToService(*discovered_device, service_uuid, &flag);
+      if (!socket_a.IsValid()) server_socket.Close();
+    });
     server_executor.Execute([&socket_b, &server_socket]() {
       socket_b = server_socket.Accept();
       if (!socket_b.IsValid()) server_socket.Close();

cpp/platform/public/wifi_lan_test.cc

@@ -64,6 +64,7 @@ TEST_P(WifiLanMediumTest, CanStartAcceptingConnectionsAndConnect) {
   std::string endpoint_info_name{kEndpointName};
   CountDownLatch found_latch(1);
   CountDownLatch accepted_latch(1);
+  CancellationFlag flag;
 
   WifiLanService* discovered_service = nullptr;
   wifi_a.StartDiscovery(
@@ -102,8 +103,7 @@ TEST_P(WifiLanMediumTest, CanStartAcceptingConnectionsAndConnect) {
   {
     SingleThreadExecutor client_executor;
     client_executor.Execute(
-        [&wifi_a, &socket_a, discovered_service, &service_id]() {
-          CancellationFlag flag;
+        [&wifi_a, &socket_a, discovered_service, &service_id, &flag]() {
           socket_a = wifi_a.Connect(*discovered_service, service_id, &flag);
         });
   }
@@ -126,6 +126,7 @@ TEST_P(WifiLanMediumTest, CanCancelConnect) {
   std::string endpoint_info_name{kEndpointName};
   CountDownLatch found_latch(1);
   CountDownLatch accepted_latch(1);
+  CancellationFlag flag(true);
 
   WifiLanService* discovered_service = nullptr;
   wifi_a.StartDiscovery(
@@ -164,9 +165,7 @@ TEST_P(WifiLanMediumTest, CanCancelConnect) {
   {
     SingleThreadExecutor client_executor;
     client_executor.Execute(
-        [&wifi_a, &socket_a, discovered_service, &service_id]() {
-          // Make it as Cancelled.
-          CancellationFlag flag(true);
+        [&wifi_a, &socket_a, discovered_service, &service_id, &flag]() {
           socket_a = wifi_a.Connect(*discovered_service, service_id, &flag);
         });
   }