From 9f7fae56b0ea23e279994f873c6a17428ede854b Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Mon, 14 Mar 2022 08:17:14 +0300
Subject: [PATCH 1/2] WriteBufferFromS3: remove unused headers

Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
---
 src/IO/WriteBufferFromS3.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/IO/WriteBufferFromS3.h b/src/IO/WriteBufferFromS3.h
index 8b89626ee18..a4fbcbcdeeb 100644
--- a/src/IO/WriteBufferFromS3.h
+++ b/src/IO/WriteBufferFromS3.h
@@ -6,6 +6,7 @@
 
 #    include <memory>
 #    include <vector>
+#    include <list>
 #    include <base/logger_useful.h>
 #    include <base/types.h>
 
@@ -14,8 +15,6 @@
 
 #    include <aws/core/utils/memory/stl/AWSStringStream.h>
 
-#    include <Common/ThreadPool.h>
-
 namespace Aws::S3
 {
 class S3Client;

From d13fc9ec44767481fa55c28d827b3a9aaf389b18 Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Mon, 14 Mar 2022 08:21:38 +0300
Subject: [PATCH 2/2] Wait for IDiskRemote thread pool properly

ASan report [1]:

    WARNING: ThreadSanitizer: heap-use-after-free (pid=611)
      Read of size 4 at 0x7b48006f0250 by thread T273:
        0 MemoryTracker::allocImpl() obj-x86_64-linux-gnu/../src/Common/MemoryTracker.cpp:105:49 (clickhouse+0xa05d0ef)
        1 MemoryTracker::allocImpl() obj-x86_64-linux-gnu/../src/Common/MemoryTracker.cpp:239:22 (clickhouse+0xa05d580)
        2 MemoryTracker::alloc(long) obj-x86_64-linux-gnu/../src/Common/MemoryTracker.cpp:246:5 (clickhouse+0xa05dea7)
        3 DB::ThreadStatus::~ThreadStatus() obj-x86_64-linux-gnu/../src/Common/ThreadStatus.cpp:154:28 (clickhouse+0xa059b01)
        4 ThreadFromGlobalPool::ThreadFromGlobalPool<>()::'lambda'()::operator()() obj-x86_64-linux-gnu/../s>

      Previous write of size 8 at 0x7b48006f0250 by main thread:
        8 std::__1::unordered_map<std::__1::basic_string<>, DB::ProcessListForUser, std::__1::hash<>, std::__1>
        9 DB::ProcessList::~ProcessList() obj-x86_64-linux-gnu/../src/Interpreters/ProcessList.h:275:7 (clickhouse+0x17b71c9c)
        10 DB::ContextSharedPart::~ContextSharedPart() obj-x86_64-linux-gnu/../src/Interpreters/Context.cpp:304:5 (clickhouse+0x17b70712)
        13 DB::SharedContextHolder::reset() obj-x86_64-linux-gnu/../src/Interpreters/Context.cpp:453:44 (clickhouse+0x17b3daf9)
        14 DB::Server::main()::$_8::o>

      Thread T273 (tid=946, running) created by thread T223 at:
        0 pthread_create <null> (clickhouse+0x9f7d49d)
        15 DB::WriteBufferFromS3::makeSinglepartUpload() obj-x86_64-linux-gnu/../src/IO/WriteBufferFromS3.cpp:332:9 (clickhouse+0x17a81b6a)

  [1]: https://s3.amazonaws.com/clickhouse-test-reports/35072/19be9c8c6433ffd1bbf8eb08975d8afbaebf3d43/stress_test__thread__actions_.html

Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
---
 src/Interpreters/Context.cpp | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/Interpreters/Context.cpp b/src/Interpreters/Context.cpp
index 6f0df4202a0..ac1bfc620b0 100644
--- a/src/Interpreters/Context.cpp
+++ b/src/Interpreters/Context.cpp
@@ -29,6 +29,7 @@
 #include <Storages/CompressionCodecSelector.h>
 #include <Storages/StorageS3Settings.h>
 #include <Disks/DiskLocal.h>
+#include <Disks/IDiskRemote.h>
 #include <TableFunctions/TableFunctionFactory.h>
 #include <Interpreters/ActionLocksManager.h>
 #include <Interpreters/ExternalLoaderXMLConfigRepository.h>
@@ -293,6 +294,17 @@ struct ContextSharedPart
 
     ~ContextSharedPart()
     {
+        /// Wait for thread pool for background writes,
+        /// since it may use per-user MemoryTracker which will be destroyed here.
+        try
+        {
+            IDiskRemote::getThreadPoolWriter().wait();
+        }
+        catch (...)
+        {
+            tryLogCurrentException(__PRETTY_FUNCTION__);
+        }
+
         try
         {
             shutdown();