Fix data race on Socket_ in TcpConnection

4eb18148c1fa8156e6abab24d2063c10a43e61c7

2 files changed, 17 insertions, 7 deletions

diff --git a/yt/yt/core/bus/tcp/connection.cpp b/yt/yt/core/bus/tcp/connection.cpp
index 7423eb8ee5..8e29b6f056 100644
--- a/yt/yt/core/bus/tcp/connection.cpp
+++ b/yt/yt/core/bus/tcp/connection.cpp
@@ -196,7 +196,7 @@ void TTcpConnection::Close()
 
     EncodedFragments_.clear();
 
-    FlushStatistics();
+    FlushStatistics(/*guarded*/ false);
 }
 
 void TTcpConnection::Start()
@@ -257,7 +257,7 @@ void TTcpConnection::RunPeriodicCheck()
         return;
     }
 
-    FlushStatistics();
+    FlushStatistics(/*guarded*/ false);
 
     auto now = NProfiling::GetCpuInstant();
 
@@ -423,7 +423,7 @@ void TTcpConnection::Open(TGuard<NThreading::TSpinLock>& guard)
     }
 
     UpdateConnectionCount(+1);
-    FlushStatistics();
+    FlushStatistics(/*guarded*/ true);
 
     // Go online and start event processing.
     auto previousPendingControl = static_cast<EPollControl>(PendingControl_.fetch_and(~static_cast<ui64>(EPollControl::Offline)));
@@ -1787,10 +1787,20 @@ void TTcpConnection::InitSocketTosLevel(TTosLevel tosLevel)
     }
 }
 
-void TTcpConnection::FlushStatistics()
+void TTcpConnection::FlushStatistics(bool guarded)
 {
-    UpdateTcpStatistics();
-    FlushBusStatistics();
+    auto doFlush = [&] {
+        UpdateTcpStatistics();
+        FlushBusStatistics();
+    };
+
+    if (guarded) {
+        doFlush();
+        return;
+    }
+
+    auto guard = Guard(Lock_);
+    doFlush();
 }
 
 template <class T, class U>
diff --git a/yt/yt/core/bus/tcp/connection.h b/yt/yt/core/bus/tcp/connection.h
index 81e2c4a3a0..5cab9e591b 100644
--- a/yt/yt/core/bus/tcp/connection.h
+++ b/yt/yt/core/bus/tcp/connection.h
@@ -359,7 +359,7 @@ private:
     void IncrementPendingOut(i64 packetSize);
     void DecrementPendingOut(i64 packetSize);
 
-    void FlushStatistics();
+    void FlushStatistics(bool guarded = false);
 
     template <class T, class U>
     i64 UpdateBusCounter(T TBusNetworkBandCounters::* field, U delta);