Merge pull request ClickHouse#79540 from ClickHouse/chesema-fix-02789

fix test 02789_reading_from_s3_with_connection_pool

Author: CheSema

src/Interpreters/AsynchronousInsertQueue.cpp

@@ -946,7 +946,7 @@ try
         throw;
     }
 
-    auto finish_entries = [&](size_t num_rows, size_t num_bytes)
+    auto finish_entries = [&](QueryPipeline && pileline_, size_t num_rows, size_t num_bytes)
     {
         for (const auto & entry : data->entries)
         {
@@ -963,7 +963,7 @@ try
         LOG_DEBUG(log, "Flushed {} rows, {} bytes for query '{}'", num_rows, num_bytes, key.query_str);
         queue_shard_flush_time_history.updateWithCurrentTime();
 
-        logQueryFinish(query_log_elem, insert_context, key.query, pipeline, /*pulling_pipeline=*/false, query_span, QueryResultCacheUsage::None, internal);
+        logQueryFinish(query_log_elem, insert_context, key.query, std::move(pileline_), /*pulling_pipeline=*/false, query_span, QueryResultCacheUsage::None, internal);
     };
 
     try
@@ -980,8 +980,8 @@ try
 
         if (chunk.getNumRows() == 0)
         {
-            finish_entries(/*num_rows=*/ 0, /*num_bytes=*/ 0);
-            pipeline.cancel();
+            pipeline.cancel(); // this just cancels the processors
+            finish_entries(std::move(pipeline), /*num_rows=*/ 0, /*num_bytes=*/ 0);
             return;
         }
 
@@ -994,7 +994,7 @@ try
         CompletedPipelineExecutor completed_executor(pipeline);
         completed_executor.execute();
 
-        finish_entries(num_rows, num_bytes);
+        finish_entries(std::move(pipeline), num_rows, num_bytes);
     }
     catch (...)
     {

src/Interpreters/executeQuery.cpp

@@ -573,7 +573,7 @@ void logQueryFinishImpl(
     QueryLogElement & elem,
     const ContextMutablePtr & context,
     const ASTPtr & query_ast,
-    const QueryPipeline & query_pipeline,
+    QueryPipeline && query_pipeline,
     bool pulling_pipeline,
     std::shared_ptr<OpenTelemetry::SpanHolder> query_span,
     QueryResultCacheUsage query_result_cache_usage,
@@ -586,16 +586,24 @@ void logQueryFinishImpl(
     QueryStatusPtr process_list_elem = context->getProcessListElement();
     if (process_list_elem)
     {
+
+        logProcessorProfile(context, query_pipeline.getProcessors());
+
+        auto result_progress = flushQueryProgress(query_pipeline, pulling_pipeline, context->getProgressCallback(), process_list_elem);
+
+        /// Reset pipeline before fetching profile counters
+        query_pipeline.reset();
+
         /// Update performance counters before logging to query_log
         CurrentThread::finalizePerformanceCounters();
 
         QueryStatusInfo info = process_list_elem->getInfo(true, settings[Setting::log_profile_events]);
         logQueryMetricLogFinish(context, internal, elem.client_info.current_query_id, time, std::make_shared<QueryStatusInfo>(info));
+
         elem.type = QueryLogElementType::QUERY_FINISH;
 
         addStatusInfoToQueryLogElement(elem, info, query_ast, context, time);
 
-        auto result_progress = flushQueryProgress(query_pipeline, pulling_pipeline, context->getProgressCallback(), process_list_elem);
         elem.result_rows = result_progress.result_rows;
         elem.result_bytes = result_progress.result_bytes;
 
@@ -622,7 +630,6 @@ void logQueryFinishImpl(
                 query_log->add(elem);
         }
 
-        logProcessorProfile(context, query_pipeline.getProcessors());
     }
 
     if (query_span)
@@ -659,14 +666,14 @@ void logQueryFinish(
     QueryLogElement & elem,
     const ContextMutablePtr & context,
     const ASTPtr & query_ast,
-    const QueryPipeline & query_pipeline,
+    QueryPipeline && query_pipeline,
     bool pulling_pipeline,
     std::shared_ptr<OpenTelemetry::SpanHolder> query_span,
     QueryResultCacheUsage query_result_cache_usage,
     bool internal)
 {
     const auto time_now = std::chrono::system_clock::now();
-    logQueryFinishImpl(elem, context, query_ast, query_pipeline, pulling_pipeline, query_span, query_result_cache_usage, internal, time_now);
+    logQueryFinishImpl(elem, context, query_ast, std::move(query_pipeline), pulling_pipeline, query_span, query_result_cache_usage, internal, time_now);
 }
 
 void logQueryException(
@@ -1643,14 +1650,14 @@ static BlockIO executeQueryImpl(
                                     execute_implicit_tcl_query,
                                     // Need to be cached, since will be changed after complete()
                                     pulling_pipeline = pipeline.pulling(),
-                                    query_span](QueryPipeline & query_pipeline, std::chrono::system_clock::time_point finish_time) mutable
+                                    query_span](QueryPipeline && query_pipeline, std::chrono::system_clock::time_point finish_time) mutable
             {
                 if (query_result_cache_usage == QueryResultCacheUsage::Write)
                     /// Trigger the actual write of the buffered query result into the query result cache. This is done explicitly to
                     /// prevent partial/garbage results in case of exceptions during query execution.
                     query_pipeline.finalizeWriteInQueryResultCache();
 
-                logQueryFinishImpl(elem, context, out_ast, query_pipeline, pulling_pipeline, query_span, query_result_cache_usage, internal, finish_time);
+                logQueryFinishImpl(elem, context, out_ast, std::move(query_pipeline), pulling_pipeline, query_span, query_result_cache_usage, internal, finish_time);
 
                 if (*implicit_txn_control)
                     execute_implicit_tcl_query(context, ASTTransactionControl::COMMIT);

src/Interpreters/executeQuery.h

@@ -93,7 +93,7 @@ void logQueryFinish(
     QueryLogElement & elem,
     const ContextMutablePtr & context,
     const ASTPtr & query_ast,
-    const QueryPipeline & query_pipeline,
+    QueryPipeline && query_pipeline,
     bool pulling_pipeline,
     std::shared_ptr<OpenTelemetry::SpanHolder> query_span,
     QueryResultCacheUsage query_result_cache_usage,

src/QueryPipeline/BlockIO.cpp

@@ -49,9 +49,9 @@ BlockIO::~BlockIO()
 void BlockIO::onFinish(std::chrono::system_clock::time_point finish_time)
 {
     if (finish_callback)
-        finish_callback(pipeline, finish_time);
-
-    pipeline.reset();
+        finish_callback(std::move(pipeline), finish_time);
+    else
+        pipeline.reset();
 }
 
 void BlockIO::onException(bool log_as_error)

src/QueryPipeline/BlockIO.h

@@ -25,7 +25,7 @@ struct BlockIO
     QueryPipeline pipeline;
 
     /// Callbacks for query logging could be set here.
-    std::function<void(QueryPipeline &, std::chrono::system_clock::time_point)> finish_callback;
+    std::function<void(QueryPipeline &&, std::chrono::system_clock::time_point)> finish_callback;
     std::function<void(bool)> exception_callback;
 
     /// When it is true, don't bother sending any non-empty blocks to the out stream

src/Storages/MaterializedView/RefreshTask.cpp

@@ -655,7 +655,7 @@ std::optional<UUID> RefreshTask::executeRefreshUnlocked(bool append, int32_t roo
             if (execution.interrupt_execution.load())
                 throw Exception(ErrorCodes::QUERY_WAS_CANCELLED, "Refresh for view {} cancelled", view_storage_id.getFullTableName());
 
-            logQueryFinish(*query_log_elem, refresh_context, refresh_query, pipeline, /*pulling_pipeline=*/false, query_span, QueryResultCacheUsage::None, /*internal=*/false);
+            logQueryFinish(*query_log_elem, refresh_context, refresh_query, std::move(pipeline), /*pulling_pipeline=*/false, query_span, QueryResultCacheUsage::None, /*internal=*/false);
             query_log_elem = std::nullopt;
             query_span = nullptr;
             process_list_entry.reset(); // otherwise it needs to be alive for logQueryException

src/Storages/StorageURL.cpp

@@ -421,7 +421,7 @@ StorageURLSource::StorageURLSource(
             if (need_only_count)
                 input_format->needOnlyCount();
 
-          builder.init(Pipe(input_format));
+            builder.init(Pipe(input_format));
 
             if (columns_description.hasDefaults())
             {

tests/queries/0_stateless/02789_reading_from_s3_with_connection_pool.reference

@@ -1,2 +1,2 @@
-1
-1
+DiskConnectionsPreserved 1
+StorageConnectionsPreserved 1

tests/queries/0_stateless/02789_reading_from_s3_with_connection_pool.sh

@@ -1,5 +1,8 @@
 #!/usr/bin/env bash
-# Tags: no-fasttest, no-random-settings, no-replicated-database, no-distributed-cache
+# Tags: no-fasttest, no-random-settings, no-replicated-database, no-distributed-cache, no-parallel-replicas
+
+# no-fasttest -- test uses s3_dick
+# no-parallel-replicas -- do not run url functions as StorageURLCluster
 
 CUR_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 # shellcheck source=../shell_config.sh
@@ -12,14 +15,14 @@ CREATE TABLE test_s3 (a UInt64, b UInt64)
 ENGINE = MergeTree ORDER BY a
 SETTINGS disk = 's3_disk', min_bytes_for_wide_part = 0;
 
-INSERT INTO test_s3 SELECT number, number FROM numbers_mt(1e7);
+INSERT INTO test_s3 SELECT number, number FROM numbers_mt(1);
 "
 
 # This (reusing connections from the pool) is not guaranteed to always happen,
 # (due to random time difference between the queries and random activity in parallel)
 # but should happen most of the time.
 
-while true
+for _ in {0..9}
 do
     query="SELECT a, b FROM test_s3"
     query_id=$(${CLICKHOUSE_CLIENT} --query "select queryID() from ($query) limit 1" 2>&1)
@@ -33,34 +36,31 @@ do
         AND query_id='$query_id';
     ")
 
-    [[ $RES -eq 1 ]] && echo "$RES" && break;
+    [[ $RES -eq 1 ]] && echo "DiskConnectionsPreserved $RES" && break;
 done
 
 
 # Test connection pool in ReadWriteBufferFromHTTP
 
-while true
+# LIMIT 1 here is important part
+# processor StorageURLSource releases the HTTP session either when all data is fully read in the last generate call or at d-tor of the processor instance
+# with LIMIT 1 the HTTP session is released at pipeline d-tor because not all the data is read from HTTP connection inside StorageURLSource processor
+# this tests covers the case when profile events have to be gathered and logged to the query_log only after pipeline is destroyed
+
+for _ in {0..9}
 do
     query_id=$(${CLICKHOUSE_CLIENT} -q "
-    create table if not exists mut (n int, m int, k int) engine=ReplicatedMergeTree('/test/02441/{database}/mut', '1') order by n;
-    set insert_keeper_fault_injection_probability=0;
-    set parallel_replicas_for_cluster_engines=0;
-    insert into mut values (1, 2, 3), (10, 20, 30);
-
-    system stop merges mut;
-    alter table mut delete where n = 10;
-
-    select queryID() from(
-        -- a funny way to wait for a MUTATE_PART to be assigned
-        select sleepEachRow(2) from url('http://localhost:8123/?param_tries={1..10}&query=' || encodeURLComponent(
-            'select 1 where ''MUTATE_PART'' not in (select type from system.replication_queue where database=''' || currentDatabase() || ''' and table=''mut'')'
-            ), 'LineAsString', 's String')
-        -- queryID() will be returned for each row, since the query above doesn't return anything we need to return a fake row
-        union all
-        select 1
-    ) limit 1 settings max_threads=1, http_make_head_request=0;
-    " 2>&1)
+    SELECT queryID() FROM(
+        SELECT *
+        FROM url(
+            'http://localhost:8123/?query=' || encodeURLComponent('select 1'),
+            'LineAsString',
+            's String')
+        ) LIMIT 1 SETTINGS max_threads=1, http_make_head_request=0;
+    ")
+
     ${CLICKHOUSE_CLIENT} --query "SYSTEM FLUSH LOGS query_log"
+
     RES=$(${CLICKHOUSE_CLIENT} -m --query "
     SELECT ProfileEvents['StorageConnectionsPreserved'] > 0
     FROM system.query_log
@@ -69,5 +69,5 @@ do
         AND query_id='$query_id';
     ")
 
-    [[ $RES -eq 1 ]] && echo "$RES" && break;
+    [[ $RES -eq 1 ]] && echo "StorageConnectionsPreserved $RES" && break;
 done