From 08229390dfa3263853662e0f39fa393f61b95a44 Mon Sep 17 00:00:00 2001
From: Yakov Olkhovskiy <yakov@clickhouse.com>
Date: Mon, 29 Jul 2024 06:17:35 +0000
Subject: [PATCH] protect socket timeouts

---
 src/Server/TCPHandler.cpp | 7 ++++++-
 src/Server/TCPHandler.h   | 3 +++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/Server/TCPHandler.cpp b/src/Server/TCPHandler.cpp
index 29568baba58..d184074729b 100644
--- a/src/Server/TCPHandler.cpp
+++ b/src/Server/TCPHandler.cpp
@@ -366,7 +366,10 @@ void TCPHandler::runImpl()
         try
         {
             /// If a user passed query-local timeouts, reset socket to initial state at the end of the query
-            SCOPE_EXIT({state.timeout_setter.reset();});
+            SCOPE_EXIT({
+                std::scoped_lock lock(in_mutex, out_mutex);
+                state.timeout_setter.reset();
+            });
 
             /** If Query - process it. If Ping or Cancel - go back to the beginning.
              *  There may come settings for a separate query that modify `query_context`.
@@ -779,6 +782,8 @@ void TCPHandler::extractConnectionSettingsFromContext(const ContextPtr & context
 
 bool TCPHandler::readDataNext()
 {
+    std::scoped_lock lock(in_mutex);
+
     Stopwatch watch(CLOCK_MONOTONIC_COARSE);
 
     /// Poll interval should not be greater than receive_timeout
diff --git a/src/Server/TCPHandler.h b/src/Server/TCPHandler.h
index 74afb5a14a5..50ef6bcf20d 100644
--- a/src/Server/TCPHandler.h
+++ b/src/Server/TCPHandler.h
@@ -229,6 +229,9 @@ private:
     /// `out_mutex` protects `out` (WriteBuffer).
     /// So it is used for method sendData(), sendProgress(), sendLogs(), etc.
     std::mutex out_mutex;
+    /// `in_mutex` protects `in` (ReadBuffer)
+    /// Used in readDataNext() and to protect socket timeout settings
+    std::mutex in_mutex;
     /// `task_callback_mutex` protects tasks callbacks.
     /// Inside these callbacks we might also change cancellation status,
     /// so it also protects cancellation status checks.