From 8e4414aaae86c58fec6306ac1805b846c09c235b Mon Sep 17 00:00:00 2001
From: packy92 <wangchao@starrocks.com>
Date: Wed, 15 Nov 2023 09:49:38 +0800
Subject: [PATCH] [BugFix] remove threadLocal after operation (#34758)

Signed-off-by: packy92 <wangchao@starrocks.com>
---
 .../src/main/java/com/starrocks/qe/StmtExecutor.java       | 7 ++++++-
 .../main/java/com/starrocks/scheduler/TaskRunExecutor.java | 2 ++
 .../optimizer/statistics/ColumnBasicStatsCacheLoader.java  | 4 ++++
 .../statistics/ColumnHistogramStatsCacheLoader.java        | 4 ++++
 .../sql/optimizer/statistics/TableStatsCacheLoader.java    | 4 ++++
 5 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java
index d9934035be89b8..488a5b73468e29 100644
--- a/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java
+++ b/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java
@@ -1093,7 +1093,12 @@ private void executeAnalyze(AnalyzeStmt analyzeStmt, AnalyzeStatus analyzeStatus
         statsConnectCtx.getSessionVariable().setStatisticCollectParallelism(
                 context.getSessionVariable().getStatisticCollectParallelism());
         statsConnectCtx.setThreadLocalInfo();
-        executeAnalyze(statsConnectCtx, analyzeStmt, analyzeStatus, db, table);
+        try {
+            executeAnalyze(statsConnectCtx, analyzeStmt, analyzeStatus, db, table);
+        } finally {
+            ConnectContext.remove();
+        }
+
     }
 
     private void executeAnalyze(ConnectContext statsConnectCtx, AnalyzeStmt analyzeStmt, AnalyzeStatus analyzeStatus,
diff --git a/fe/fe-core/src/main/java/com/starrocks/scheduler/TaskRunExecutor.java b/fe/fe-core/src/main/java/com/starrocks/scheduler/TaskRunExecutor.java
index 92721593937e29..3e881383453e5c 100644
--- a/fe/fe-core/src/main/java/com/starrocks/scheduler/TaskRunExecutor.java
+++ b/fe/fe-core/src/main/java/com/starrocks/scheduler/TaskRunExecutor.java
@@ -15,6 +15,7 @@
 
 package com.starrocks.scheduler;
 
+import com.starrocks.qe.ConnectContext;
 import com.starrocks.scheduler.persist.TaskRunStatus;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -56,6 +57,7 @@ public void executeTaskRun(TaskRun taskRun) {
                 status.setErrorCode(-1);
                 status.setErrorMessage(ex.toString());
             } finally {
+                ConnectContext.remove();
                 status.setFinishTime(System.currentTimeMillis());
             }
             return status.getState();
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnBasicStatsCacheLoader.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnBasicStatsCacheLoader.java
index f9191e79e67c60..902a272b5d4524 100644
--- a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnBasicStatsCacheLoader.java
+++ b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnBasicStatsCacheLoader.java
@@ -68,6 +68,8 @@ public class ColumnBasicStatsCacheLoader implements AsyncCacheLoader<ColumnStats
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
@@ -105,6 +107,8 @@ public class ColumnBasicStatsCacheLoader implements AsyncCacheLoader<ColumnStats
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnHistogramStatsCacheLoader.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnHistogramStatsCacheLoader.java
index d1930f86595084..bfaf9743dfb425 100644
--- a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnHistogramStatsCacheLoader.java
+++ b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnHistogramStatsCacheLoader.java
@@ -75,6 +75,8 @@ CompletableFuture<Optional<Histogram>> asyncLoad(@NonNull ColumnStatsCacheKey ca
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
@@ -107,6 +109,8 @@ CompletableFuture<Optional<Histogram>> asyncLoad(@NonNull ColumnStatsCacheKey ca
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/TableStatsCacheLoader.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/TableStatsCacheLoader.java
index 6d6394092f9c51..fcd11f2e8365ab 100644
--- a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/TableStatsCacheLoader.java
+++ b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/statistics/TableStatsCacheLoader.java
@@ -50,6 +50,8 @@ public class TableStatsCacheLoader implements AsyncCacheLoader<TableStatsCacheKe
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
@@ -84,6 +86,8 @@ public class TableStatsCacheLoader implements AsyncCacheLoader<TableStatsCacheKe
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }