From 86119dbc3f56a78234eaa67ea40760bb0e30e7fa Mon Sep 17 00:00:00 2001
From: Sema Checherinda <Sema.Checherinda@clickhouse.com>
Date: Fri, 17 Nov 2023 12:26:50 +0100
Subject: [PATCH] fix data race

---
 src/Storages/MergeTree/DataPartsExchange.cpp           | 10 ++++++++--
 .../02916_replication_protocol_wait_for_part.sql       |  4 ++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/Storages/MergeTree/DataPartsExchange.cpp b/src/Storages/MergeTree/DataPartsExchange.cpp
index c39263a0b73..0192fb1868b 100644
--- a/src/Storages/MergeTree/DataPartsExchange.cpp
+++ b/src/Storages/MergeTree/DataPartsExchange.cpp
@@ -349,7 +349,7 @@ MergeTreeData::DataPart::Checksums Service::sendPartFromDisk(
     return data_checksums;
 }
 
-bool wait_loop(UInt32 wait_timeout_ms, std::function<bool()> pred)
+bool wait_loop(UInt32 wait_timeout_ms, const std::function<bool()> & pred)
 {
     static const UInt32 loop_delay_ms = 5;
 
@@ -360,6 +360,7 @@ bool wait_loop(UInt32 wait_timeout_ms, std::function<bool()> pred)
         return true;
 
     Stopwatch timer;
+    sleepForMilliseconds(loop_delay_ms);
     while (!pred() && timer.elapsedMilliseconds() < wait_timeout_ms)
     {
         sleepForMilliseconds(loop_delay_ms);
@@ -387,8 +388,13 @@ MergeTreeData::DataPartPtr Service::findPart(const String & name)
     /// do not expose PreActive parts for zero-copy
 
     static const UInt32 wait_timeout_ms = 1000;
-    bool pred_result = wait_loop(wait_timeout_ms, [&] () { return part->getState() != MergeTreeDataPartState::PreActive; });
+    auto pred = [&] ()
+    {
+        auto lock = data.lockParts();
+        return part->getState() != MergeTreeDataPartState::PreActive;
+    };
 
+    bool pred_result = wait_loop(wait_timeout_ms, pred);
     if (!pred_result)
         throw Exception(
                 ErrorCodes::ABORTED,
diff --git a/tests/queries/0_stateless/02916_replication_protocol_wait_for_part.sql b/tests/queries/0_stateless/02916_replication_protocol_wait_for_part.sql
index 97ef33f96e8..010e29a34e8 100644
--- a/tests/queries/0_stateless/02916_replication_protocol_wait_for_part.sql
+++ b/tests/queries/0_stateless/02916_replication_protocol_wait_for_part.sql
@@ -22,5 +22,5 @@ insert into tableIn values(2);
 system sync replica tableOut;
 select count() from tableOut;
 
-drop table tableIn
-drop table tableOut
+drop table tableIn;
+drop table tableOut;