[Enhancement] Support profile for only big query (StarRocks#33825)

Signed-off-by: liuyehcf <[email protected]>
(cherry picked from commit 0e2d056)

be/src/exec/pipeline/fragment_executor.cpp

@@ -131,6 +131,9 @@ Status FragmentExecutor::_prepare_query_ctx(ExecEnv* exec_env, const UnifiedExec
     if (query_options.__isset.enable_profile && query_options.enable_profile) {
         _query_ctx->set_report_profile();
     }
+    if (query_options.__isset.big_query_profile_second_threshold) {
+        _query_ctx->set_big_query_profile_threshold(query_options.big_query_profile_second_threshold);
+    }
     if (query_options.__isset.pipeline_profile_level) {
         _query_ctx->set_profile_level(query_options.pipeline_profile_level);
     }

be/src/exec/pipeline/query_context.h

@@ -73,7 +73,18 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
                 duration_cast<milliseconds>(steady_clock::now().time_since_epoch() + _query_expire_seconds).count();
     }
     void set_report_profile() { _is_report_profile = true; }
-    bool is_report_profile() { return _is_report_profile; }
+    bool is_report_profile() {
+        if (_is_report_profile) {
+            return true;
+        }
+        if (_big_query_profile_threshold_ns <= 0) {
+            return false;
+        }
+        return MonotonicNanos() - _query_begin_time > _big_query_profile_threshold_ns;
+    }
+    void set_big_query_profile_threshold(int64_t big_query_profile_threshold_s) {
+        _big_query_profile_threshold_ns = 1'000'000'000L * big_query_profile_threshold_s;
+    }
     void set_profile_level(const TPipelineProfileLevel::type& profile_level) { _profile_level = profile_level; }
     const TPipelineProfileLevel::type& profile_level() { return _profile_level; }
 
@@ -166,6 +177,7 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
     std::once_flag _init_mem_tracker_once;
     std::shared_ptr<RuntimeProfile> _profile;
     bool _is_report_profile = false;
+    int64_t _big_query_profile_threshold_ns = 0;
     TPipelineProfileLevel::type _profile_level;
     std::shared_ptr<MemTracker> _mem_tracker;
     ObjectPool _object_pool;

fe/fe-core/src/main/java/com/starrocks/load/loadv2/LoadLoadingTask.java

@@ -174,7 +174,7 @@ private void executeOnce() throws Exception {
             long beginTimeInNanoSecond = TimeUtils.getStartTime();
             actualExecute(curCoordinator);
 
-            if (context.getSessionVariable().isEnableProfile()) {
+            if (context.isProfileEnabled()) {
                 RuntimeProfile profile = new RuntimeProfile("Load");
                 RuntimeProfile summaryProfile = new RuntimeProfile("Summary");
                 summaryProfile.addInfoString(ProfileManager.QUERY_ID, DebugUtil.printId(context.getExecutionId()));

fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java

@@ -496,6 +496,20 @@ public void setLastQueryId(UUID queryId) {
         this.lastQueryId = queryId;
     }
 
+    public boolean isProfileEnabled() {
+        if (sessionVariable == null) {
+            return false;
+        }
+        if (sessionVariable.isEnableProfile()) {
+            return true;
+        }
+        if (!sessionVariable.isEnableBigQueryProfile()) {
+            return false;
+        }
+        return System.currentTimeMillis() - getStartTime() >
+                1000L * sessionVariable.getBigQueryProfileSecondThreshold();
+    }
+
     public byte[] getAuthDataSalt() {
         return authDataSalt;
     }

fe/fe-core/src/main/java/com/starrocks/qe/Coordinator.java

@@ -286,7 +286,8 @@ public Coordinator(ConnectContext context, List<PlanFragment> fragments, List<Sc
         if (context.getLastQueryId() != null) {
             this.queryGlobals.setLast_query_id(context.getLastQueryId().toString());
         }
-        this.needReport = context.getSessionVariable().isEnableProfile();
+        this.needReport = context.getSessionVariable().isEnableProfile() ||
+                context.getSessionVariable().isEnableBigQueryProfile();
         this.preferComputeNode = context.getSessionVariable().isPreferComputeNode();
         this.useComputeNodeNumber = context.getSessionVariable().getUseComputeNodes();
         this.nextInstanceId = new TUniqueId();
@@ -1633,7 +1634,7 @@ private void cancelInternal(PPlanFragmentCancelReason cancelReason) {
         cancelRemoteFragmentsAsync(cancelReason);
         if (profileDoneSignal != null && cancelReason != PPlanFragmentCancelReason.LIMIT_REACH) {
             // count down to zero to notify all objects waiting for this
-            if (!connectContext.getSessionVariable().isEnableProfile()) {
+            if (!connectContext.isProfileEnabled()) {
                 profileDoneSignal.countDownToZero(new Status());
                 LOG.info("unfinished instance: {}",
                         profileDoneSignal.getLeftMarks().stream().map(e -> DebugUtil.printId(e.getKey())).toArray());
@@ -2615,7 +2616,7 @@ public boolean join(int timeoutS) {
     public void mergeIsomorphicProfiles() {
         SessionVariable sessionVariable = connectContext.getSessionVariable();
 
-        if (!sessionVariable.isEnableProfile()) {
+        if (!connectContext.isProfileEnabled()) {
             return;
         }
 

fe/fe-core/src/main/java/com/starrocks/qe/SessionVariable.java

@@ -332,6 +332,8 @@ public class SessionVariable implements Serializable, Writable, Cloneable {
     public static final String ENABLE_MATERIALIZED_VIEW_SINGLE_TABLE_VIEW_DELTA_REWRITE =
             "enable_materialized_view_single_table_view_delta_rewrite";
 
+    public static final String BIG_QUERY_PROFILE_SECOND_THRESHOLD = "big_query_profile_second_threshold";
+
     public static final String ENABLE_PRUNE_COMPLEX_TYPES = "enable_prune_complex_types";
 
     public static final String GROUP_CONCAT_MAX_LEN = "group_concat_max_len";
@@ -615,6 +617,9 @@ public class SessionVariable implements Serializable, Writable, Cloneable {
     @VariableMgr.VarAttr(name = PIPELINE_PROFILE_LEVEL)
     private int pipelineProfileLevel = 1;
 
+    @VariableMgr.VarAttr(name = BIG_QUERY_PROFILE_SECOND_THRESHOLD)
+    private int bigQueryProfileSecondThreshold = 0;
+
     @VariableMgr.VarAttr(name = RESOURCE_GROUP_ID, alias = RESOURCE_GROUP_ID_V2,
             show = RESOURCE_GROUP_ID_V2, flag = VariableMgr.INVISIBLE)
     private int resourceGroupId = 0;
@@ -1076,6 +1081,14 @@ public void setEnableProfile(boolean enableProfile) {
         this.enableProfile = enableProfile;
     }
 
+    public boolean isEnableBigQueryProfile() {
+        return bigQueryProfileSecondThreshold > 0;
+    }
+
+    public int getBigQueryProfileSecondThreshold() {
+        return bigQueryProfileSecondThreshold;
+    }
+
     public int getWaitTimeoutS() {
         return waitTimeout;
     }
@@ -1737,6 +1750,7 @@ public TQueryOptions toThrift() {
         tResult.setQuery_delivery_timeout(Math.min(Integer.MAX_VALUE / 1000, queryDeliveryTimeoutS));
         tResult.setEnable_profile(enableProfile);
         tResult.setCodegen_level(0);
+        tResult.setBig_query_profile_second_threshold(bigQueryProfileSecondThreshold);
         tResult.setBatch_size(chunkSize);
         tResult.setDisable_stream_preaggregations(disableStreamPreaggregations);
         tResult.setLoad_mem_limit(loadMemLimit);

fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java

@@ -461,7 +461,7 @@ public void execute() throws Exception {
                             throw e;
                         }
                     } finally {
-                        if (!needRetry && context.getSessionVariable().isEnableProfile()) {
+                        if (!needRetry && context.isProfileEnabled()) {
                             writeProfile(beginTimeInNanoSecond);
                         }
                         QeProcessorImpl.INSTANCE.unregisterQuery(context.getExecutionId());
@@ -579,7 +579,7 @@ private void handleCreateTableAsSelectStmt(long beginTimeInNanoSecond) throws Ex
             InsertStmt insertStmt = createTableAsSelectStmt.getInsertStmt();
             ExecPlan execPlan = new StatementPlanner().plan(insertStmt, context);
             handleDMLStmt(execPlan, ((CreateTableAsSelectStmt) parsedStmt).getInsertStmt());
-            if (context.getSessionVariable().isEnableProfile()) {
+            if (context.isProfileEnabled()) {
                 writeProfile(beginTimeInNanoSecond);
             }
             if (context.getState().getStateType() == MysqlStateType.ERR) {

gensrc/thrift/InternalService.thrift

@@ -231,6 +231,8 @@ struct TQueryOptions {
   93: optional i32 connector_io_tasks_slow_io_latency_ms = 50;
   94: optional double scan_use_query_mem_ratio = 0.25;
   95: optional double connector_scan_use_query_mem_ratio = 0.3;  
+
+  109: optional i64 big_query_profile_second_threshold;
 }