Merge pull request #320 from wudidapaopao/yuxiaozhe-dev

Fix the potential memory leak issue caused by query result memory not being properly released

Author: auxten

src/Client/ClientBase.cpp

@@ -523,7 +523,14 @@ try
         else
         {
             // The underlying vector<char> will be freed by `free_result` from Python side.
-            query_result_memory = new std::vector<char>(4096);
+            // When the input query contains multiple statements, such as "SELECT 1; SELECT 2". When executing "SELECT 2",
+            // the query_result_memory corresponding to SELECT 1 has not been "stolen" by function stealQueryOutputVector,
+            // so it remains non-empty.
+            if (query_result_memory)
+                query_result_memory->resize(4096);
+            else
+                query_result_memory = new std::vector<char>(4096);
+
             query_result_buf = std::make_shared<WriteBufferFromVector<std::vector<char>>>(*query_result_memory);
 
             out_buf = query_result_buf.get();

src/Client/ClientBase.h

@@ -107,13 +107,13 @@ class ClientBase
         return query_result_memory;
     }
 
-    /// Steals and returns the query output vector, replacing it with a new one
+    /// Steals and returns the query output vector.
+    /// Afterward, the content of query_result_memory is released by the Python side.
     std::vector<char> * stealQueryOutputVector()
     {
         auto * result = query_result_memory;
-        query_result_memory = new std::vector<char>(4096);
-        // WriteBufferFromVector takes a reference to the vector but doesn't own it
-        query_result_buf = std::make_shared<WriteBufferFromVector<std::vector<char>>>(*query_result_memory);
+        query_result_memory = nullptr;
+        query_result_buf.reset();
         return result;
     }
 

src/Core/ServerSettings.h

@@ -126,7 +126,7 @@ namespace DB
     M(UInt64, background_fetches_pool_size, 1, "The maximum number of threads that will be used for fetching data parts from another replica for *MergeTree-engine tables in a background.", 0) \
     M(UInt64, background_common_pool_size, 1, "The maximum number of threads that will be used for performing a variety of operations (mostly garbage collection) for *MergeTree-engine tables in a background.", 0) \
     M(UInt64, background_buffer_flush_schedule_pool_size, 2, "The maximum number of threads that will be used for performing flush operations for Buffer-engine tables in a background.", 0) \
-    M(UInt64, background_schedule_pool_size, 1, "The maximum number of threads that will be used for constantly executing some lightweight periodic operations.", 0) \
+    M(UInt64, background_schedule_pool_size, 2, "The maximum number of threads that will be used for constantly executing some lightweight periodic operations.", 0) \
     M(UInt64, background_message_broker_schedule_pool_size, 0, "The maximum number of threads that will be used for executing background operations for message streaming.", 0) \
     M(UInt64, background_distributed_schedule_pool_size, 0, "The maximum number of threads that will be used for executing distributed sends.", 0) \
     M(UInt64, tables_loader_foreground_pool_size, 0, "The maximum number of threads that will be used for foreground (that is being waited for by a query) loading of tables. Also used for synchronous loading of tables before the server start. Zero means use all CPUs.", 0) \