From 033d88ca246b6a445107293bffb991cffe87cb13 Mon Sep 17 00:00:00 2001
From: Francis Tsui <ftsui@google.com>
Date: Thu, 11 Apr 2024 12:40:34 -0700
Subject: [PATCH] TSAN error fixes.

PiperOrigin-RevId: 623908987
---
 sharing/BUILD                     |  3 ++-
 sharing/fake_nearby_connection.cc | 31 +++++++++++++++++++++++--------
 sharing/fake_nearby_connection.h  | 22 +++++++++++++++-------
 3 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/sharing/BUILD b/sharing/BUILD
index 1e9fbbdc..82cb4f51 100644
--- a/sharing/BUILD
+++ b/sharing/BUILD
@@ -244,10 +244,11 @@ cc_library(
         "//sharing/internal/public:logging",
         "//sharing/local_device_data",
         "//sharing/proto:enums_cc_proto",
-        "//sharing/proto:share_cc_proto",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
     ],
 )
 
diff --git a/sharing/fake_nearby_connection.cc b/sharing/fake_nearby_connection.cc
index 7e96fb39..67adc281 100644
--- a/sharing/fake_nearby_connection.cc
+++ b/sharing/fake_nearby_connection.cc
@@ -22,6 +22,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/synchronization/mutex.h"
 #include "sharing/internal/public/logging.h"
 #include "sharing/nearby_connection.h"
 
@@ -32,12 +33,16 @@ FakeNearbyConnection::~FakeNearbyConnection() = default;
 
 void FakeNearbyConnection::Read(ReadCallback callback) {
   NL_DCHECK(!closed_);
-  callback_ = std::move(callback);
+  {
+    absl::MutexLock lock(&read_mutex_);
+    callback_ = std::move(callback);
+  }
   MaybeRunCallback();
 }
 
 void FakeNearbyConnection::Write(std::vector<uint8_t> bytes) {
   NL_DCHECK(!closed_);
+  absl::MutexLock lock(&write_mutex_);
   write_data_.push(std::move(bytes));
 }
 
@@ -49,6 +54,7 @@ void FakeNearbyConnection::Close() {
     std::move(disconnect_listener_)();
   }
 
+  absl::MutexLock lock(&read_mutex_);
   if (callback_) {
     has_read_callback_been_run_ = true;
     auto callback = std::move(callback_);
@@ -65,11 +71,15 @@ void FakeNearbyConnection::SetDisconnectionListener(
 
 void FakeNearbyConnection::AppendReadableData(std::vector<uint8_t> bytes) {
   NL_DCHECK(!closed_);
-  read_data_.push(std::move(bytes));
+  {
+    absl::MutexLock lock(&read_mutex_);
+    read_data_.push(std::move(bytes));
+  }
   MaybeRunCallback();
 }
 
 std::vector<uint8_t> FakeNearbyConnection::GetWrittenData() {
+  absl::MutexLock lock(&write_mutex_);
   if (write_data_.empty()) return {};
 
   std::vector<uint8_t> bytes = std::move(write_data_.front());
@@ -81,12 +91,17 @@ bool FakeNearbyConnection::IsClosed() { return closed_; }
 
 void FakeNearbyConnection::MaybeRunCallback() {
   NL_DCHECK(!closed_);
-  if (!callback_ || read_data_.empty()) return;
-  auto item = std::move(read_data_.front());
-  read_data_.pop();
-  has_read_callback_been_run_ = true;
-  auto callback = std::move(callback_);
-  callback_ = nullptr;
+  std::vector<uint8_t> item;
+  ReadCallback callback;
+  {
+    absl::MutexLock lock(&read_mutex_);
+    if (!callback_ || read_data_.empty()) return;
+    item = std::move(read_data_.front());
+    read_data_.pop();
+    callback = std::move(callback_);
+    callback_ = nullptr;
+    has_read_callback_been_run_ = true;
+  }
   callback(std::move(item));
 }
 
diff --git a/sharing/fake_nearby_connection.h b/sharing/fake_nearby_connection.h
index c2a4e342..aecbaa05 100644
--- a/sharing/fake_nearby_connection.h
+++ b/sharing/fake_nearby_connection.h
@@ -21,6 +21,8 @@
 #include <queue>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
 #include "sharing/nearby_connection.h"
 
 namespace nearby {
@@ -37,21 +39,27 @@ class FakeNearbyConnection : public NearbyConnection {
   void Close() override;
   void SetDisconnectionListener(std::function<void()> listener) override;
 
-  void AppendReadableData(std::vector<uint8_t> bytes);
+  void AppendReadableData(std::vector<uint8_t> bytes)
+      ABSL_LOCKS_EXCLUDED(read_mutex_);
   std::vector<uint8_t> GetWrittenData();
 
   bool IsClosed();
-  bool has_read_callback_been_run() { return has_read_callback_been_run_; }
+  bool has_read_callback_been_run() {
+    absl::MutexLock lock(&read_mutex_);
+    return has_read_callback_been_run_;
+  }
 
  private:
-  void MaybeRunCallback();
+  void MaybeRunCallback() ABSL_LOCKS_EXCLUDED(read_mutex_);
 
   bool closed_ = false;
-  bool has_read_callback_been_run_ = false;
 
-  ReadCallback callback_;
-  std::queue<std::vector<uint8_t>> read_data_;
-  std::queue<std::vector<uint8_t>> write_data_;
+  absl::Mutex read_mutex_;
+  bool has_read_callback_been_run_ ABSL_GUARDED_BY(read_mutex_) = false;
+  ReadCallback callback_ ABSL_GUARDED_BY(read_mutex_);
+  std::queue<std::vector<uint8_t>> read_data_ ABSL_GUARDED_BY(read_mutex_);
+  absl::Mutex write_mutex_;
+  std::queue<std::vector<uint8_t>> write_data_ ABSL_GUARDED_BY(write_mutex_);
   std::function<void()> disconnect_listener_;
 };