minor code improvements around ThreadStatus

src/Common/CurrentThread.h

@@ -63,9 +63,6 @@ public:
     /// Call from master thread as soon as possible (e.g. when thread accepted connection)
     static void initializeQuery();
 
-    /// Sets query_context for current thread group
-    static void attachQueryContext(Context & query_context);
-
     /// You must call one of these methods when create a query child thread:
     /// Add current thread to a group associated with the thread group
     static void attachTo(const ThreadGroupStatusPtr & thread_group);
@@ -99,6 +96,10 @@ public:
 
 private:
     static void defaultThreadDeleter();
+
+    /// Sets query_context for current thread group
+    /// Can by used only through QueryScope
+    static void attachQueryContext(Context & query_context);
 };
 
 }

src/Common/ThreadStatus.cpp

@@ -99,6 +99,9 @@ ThreadStatus::~ThreadStatus()
         /// We've already allocated a little bit more than the limit and cannot track it in the thread memory tracker or its parent.
     }
 
+    /// It may cause segfault if query_context was destroyed, but was not detached
+    assert((!query_context && query_id.empty()) || (query_id == query_context->getCurrentQueryId()));
+
     if (deleter)
         deleter();
     current_thread = nullptr;

src/Common/ThreadStatus.h

@@ -201,7 +201,7 @@ public:
     void setFatalErrorCallback(std::function<void()> callback);
     void onFatalError();
 
-    /// Sets query context for current thread and its thread group
+    /// Sets query context for current master thread and its thread group
     /// NOTE: query_context have to be alive until detachQuery() is called
     void attachQueryContext(Context & query_context);
 

src/Interpreters/DDLWorker.cpp

@@ -609,12 +609,14 @@ bool DDLWorker::tryExecuteQuery(const String & query, const DDLTask & task, Exec
     ReadBufferFromString istr(query_to_execute);
     String dummy_string;
     WriteBufferFromString ostr(dummy_string);
+    std::optional<CurrentThread::QueryScope> query_scope;
 
     try
     {
         auto current_context = std::make_unique<Context>(context);
         current_context->getClientInfo().query_kind = ClientInfo::QueryKind::SECONDARY_QUERY;
         current_context->setCurrentQueryId(""); // generate random query_id
+        query_scope.emplace(*current_context);
         executeQuery(istr, ostr, false, *current_context, {});
     }
     catch (...)
@@ -631,20 +633,6 @@ bool DDLWorker::tryExecuteQuery(const String & query, const DDLTask & task, Exec
     return true;
 }
 
-void DDLWorker::attachToThreadGroup()
-{
-    if (thread_group)
-    {
-        /// Put all threads to one thread pool
-        CurrentThread::attachToIfDetached(thread_group);
-    }
-    else
-    {
-        CurrentThread::initializeQuery();
-        thread_group = CurrentThread::getGroup();
-    }
-}
-
 
 void DDLWorker::enqueueTask(DDLTaskPtr task_ptr)
 {
@@ -1124,8 +1112,6 @@ void DDLWorker::runMainThread()
     {
         try
         {
-            attachToThreadGroup();
-
             cleanup_event->set();
             scheduleTasks();
 

src/Interpreters/DDLWorker.h

@@ -161,8 +161,6 @@ private:
     void runMainThread();
     void runCleanupThread();
 
-    void attachToThreadGroup();
-
 private:
     Context context;
     Poco::Logger * log;
@@ -195,8 +193,6 @@ private:
     /// How many tasks could be in the queue
     size_t max_tasks_in_queue = 1000;
 
-    ThreadGroupStatusPtr thread_group;
-
     std::atomic<UInt64> max_id = 0;
 
     friend class DDLQueryStatusInputStream;

src/Interpreters/ThreadStatusExt.cpp

@@ -500,6 +500,8 @@ CurrentThread::QueryScope::QueryScope(Context & query_context)
 {
     CurrentThread::initializeQuery();
     CurrentThread::attachQueryContext(query_context);
+    if (!query_context.hasQueryContext())
+        query_context.makeQueryContext();
 }
 
 void CurrentThread::QueryScope::logPeakMemoryUsage()

src/Interpreters/executeQuery.cpp

@@ -326,13 +326,8 @@ static std::tuple<ASTPtr, BlockIO> executeQueryImpl(
 {
     const auto current_time = std::chrono::system_clock::now();
 
-    /// If we already executing query and it requires to execute internal query, than
-    /// don't replace thread context with given (it can be temporary). Otherwise, attach context to thread.
-    if (!internal)
-    {
-        context.makeQueryContext();
-        CurrentThread::attachQueryContext(context);
-    }
+    assert(internal || CurrentThread::get().getQueryContext());
+    assert(internal || CurrentThread::get().getQueryContext()->getCurrentQueryId() == CurrentThread::getQueryId());
 
     const Settings & settings = context.getSettingsRef();
 

src/Server/MySQLHandler.cpp

@@ -24,6 +24,7 @@
 #include <regex>
 #include <Access/User.h>
 #include <Access/AccessControlManager.h>
+#include <Common/setThreadName.h>
 
 #if !defined(ARCADIA_BUILD)
 #    include <Common/config_version.h>
@@ -86,6 +87,8 @@ MySQLHandler::MySQLHandler(IServer & server_, const Poco::Net::StreamSocket & so
 
 void MySQLHandler::run()
 {
+    setThreadName("MySQLHandler");
+    ThreadStatus thread_status;
     connection_context.makeSessionContext();
     connection_context.getClientInfo().interface = ClientInfo::Interface::MYSQL;
     connection_context.setDefaultFormat("MySQLWire");
@@ -339,8 +342,9 @@ void MySQLHandler::comQuery(ReadBuffer & payload)
 
             affected_rows += progress.written_rows;
         });
+        CurrentThread::QueryScope query_scope{query_context};
 
-        executeQuery(should_replace ? replacement : payload, *out, true, query_context,
+        executeQuery(should_replace ? replacement : payload, *out, false, query_context,
             [&with_output](const String &, const String &, const String &, const String &)
             {
                 with_output = true;

src/Server/PostgreSQLHandler.cpp

@@ -5,6 +5,7 @@
 #include <Interpreters/executeQuery.h>
 #include "PostgreSQLHandler.h"
 #include <Parsers/parseQuery.h>
+#include <Common/setThreadName.h>
 #include <random>
 
 #if !defined(ARCADIA_BUILD)
@@ -49,6 +50,8 @@ void PostgreSQLHandler::changeIO(Poco::Net::StreamSocket & socket)
 
 void PostgreSQLHandler::run()
 {
+    setThreadName("PostgresHandler");
+    ThreadStatus thread_status;
     connection_context.makeSessionContext();
     connection_context.getClientInfo().interface = ClientInfo::Interface::POSTGRESQL;
     connection_context.setDefaultFormat("PostgreSQLWire");
@@ -273,8 +276,10 @@ void PostgreSQLHandler::processQuery()
 
         for (const auto & spl_query : queries)
         {
+            /// FIXME why do we execute all queries in a single connection context?
+            CurrentThread::QueryScope query_scope{connection_context};
             ReadBufferFromString read_buf(spl_query);
-            executeQuery(read_buf, *out, true, connection_context, {});
+            executeQuery(read_buf, *out, false, connection_context, {});
 
             PostgreSQLProtocol::Messaging::CommandComplete::Command command =
                 PostgreSQLProtocol::Messaging::CommandComplete::classifyQuery(spl_query);

src/Storages/StorageReplicatedMergeTree.cpp

@@ -3680,7 +3680,7 @@ void StorageReplicatedMergeTree::shutdown()
 
     /// We clear all old parts after stopping all background operations. It's
     /// important, because background operations can produce temporary parts
-    /// which will remove themselves in their descrutors. If so, we may have
+    /// which will remove themselves in their destrutors. If so, we may have
     /// race condition between our remove call and background process.
     clearOldPartsFromFilesystem(true);
 }