From 844a17c9fbbe621d16a90c830e9a527b1f601bec Mon Sep 17 00:00:00 2001
From: Francis Tsui <ftsui@google.com>
Date: Thu, 19 Dec 2024 15:05:23 -0800
Subject: [PATCH] Clean outgoing share targets when last send surface is
 unregistered.

PiperOrigin-RevId: 708048158
---
 sharing/nearby_sharing_service_impl.cc | 35 +++++++++++++++-----------
 sharing/nearby_sharing_service_impl.h  |  2 +-
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/sharing/nearby_sharing_service_impl.cc b/sharing/nearby_sharing_service_impl.cc
index c0afcab6..6ef2d10e 100644
--- a/sharing/nearby_sharing_service_impl.cc
+++ b/sharing/nearby_sharing_service_impl.cc
@@ -510,14 +510,14 @@ void NearbySharingServiceImpl::RegisterSendSurface(
 
         // Let newly registered send surface catch up with discovered share
         // targets from current scanning session.
-        if (is_scanning_) {
+        // if (is_scanning_) {
           for (const auto& item : outgoing_share_target_map_) {
             LOG(INFO) << "Reporting discovered target "
                       << item.second.ToString()
                       << " when registering send surface";
             wrapped_callback.OnShareTargetDiscovered(item.second);
           }
-        }
+        // }
 
         // Set Share Start time for Foreground Send Surfaces
         if (state == SendSurfaceState::kForeground) {
@@ -868,7 +868,9 @@ void NearbySharingServiceImpl::Reject(
 
         RunOnNearbySharingServiceThreadDelayed(
             "incoming_rejection_delay", kIncomingRejectionDelay,
-            [this, share_target_id]() { CloseConnection(share_target_id); });
+            [this, endpoint_id = session->endpoint_id()]() {
+              CloseConnection(endpoint_id);
+            });
         // kRejected status already sent below, no need to send on disconnect.
         session->set_disconnect_status(TransferMetadata::Status::kUnknown);
 
@@ -957,9 +959,9 @@ void NearbySharingServiceImpl::DoCancel(
 
       RunOnNearbySharingServiceThreadDelayed(
           "initiator_cancel_delay", kInitiatorCancelDelay,
-          [this, share_target_id]() {
+          [this, endpoint_id = session->endpoint_id()]() {
             LOG(INFO) << "Close connection after cancellation delay.";
-            CloseConnection(share_target_id);
+            CloseConnection(endpoint_id);
           });
 
       session->WriteCancelFrame();
@@ -1109,6 +1111,17 @@ NearbySharingServiceImpl::InternalUnregisterSendSurface(
           share_target, attachment_container, transfer_metadata);
     }
   }
+  if (foreground_send_surface_map_.empty() &&
+      background_send_surface_map_.empty()) {
+    LOG(INFO) << __func__ << ": Last send surface has been unregistered";
+    // Clear outgoing_share_targets, outgoing_share_sessions and
+    // discovery_cache.
+    while (!outgoing_share_target_map_.empty()) {
+      RemoveOutgoingShareTargetWithEndpointId(
+          outgoing_share_target_map_.begin()->first);
+    }
+    discovery_cache_.clear();
+  }
 
   VLOG(1) << __func__ << ": A SendSurface has been unregistered: "
           << SendSurfaceStateToString(state);
@@ -2489,7 +2502,7 @@ void NearbySharingServiceImpl::Fail(IncomingShareSession& session,
   RunOnNearbySharingServiceThreadDelayed(
       "incoming_rejection_delay", kIncomingRejectionDelay,
       absl::bind_front(&NearbySharingServiceImpl::CloseConnection, this,
-                       session.share_target().id));
+                       session.endpoint_id()));
 
   session.SendFailureResponse(status);
 }
@@ -2648,14 +2661,8 @@ void NearbySharingServiceImpl::OnOutgoingTransferUpdate(
   }
 }
 
-void NearbySharingServiceImpl::CloseConnection(int64_t share_target_id) {
-  ShareSession* session = GetShareSession(share_target_id);
-  if (session != nullptr && session->IsConnected()) {
-    session->Disconnect();
-    return;
-  }
-  LOG(WARNING) << __func__ << ": Invalid connection for target - "
-               << share_target_id;
+void NearbySharingServiceImpl::CloseConnection(absl::string_view endpoint_id) {
+  nearby_connections_manager_->Disconnect(endpoint_id);
 }
 
 void NearbySharingServiceImpl::OnIncomingDecryptedCertificate(
diff --git a/sharing/nearby_sharing_service_impl.h b/sharing/nearby_sharing_service_impl.h
index f723568a..dba4ae9e 100644
--- a/sharing/nearby_sharing_service_impl.h
+++ b/sharing/nearby_sharing_service_impl.h
@@ -327,7 +327,7 @@ class NearbySharingServiceImpl
                                 const TransferMetadata& metadata);
   void OnOutgoingTransferUpdate(OutgoingShareSession& session,
                                 const TransferMetadata& metadata);
-  void CloseConnection(int64_t share_target_id);
+  void CloseConnection(absl::string_view endpoint_id);
   void OnIncomingDecryptedCertificate(
       absl::string_view endpoint_id, const Advertisement& advertisement,
       int64_t placeholder_share_target_id,