Merge pull request #35257 from azat/fix-s3-pool-shutdown

Wait for IDiskRemote thread pool properly

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;

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();