[BugFix] remove threadLocal after operation (StarRocks#34758)

Signed-off-by: packy92 <wangchao@starrocks.com>
packy92 · Nov 18, 2023 · 8e4414a · 8e4414a
1 parent b1c0f30
commit 8e4414a
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 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
@@ -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
@@ -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/...ore/src/main/java/com/starrocks/sql/optimizer/statistics/ColumnBasicStatsCacheLoader.java b/...ore/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/...src/main/java/com/starrocks/sql/optimizer/statistics/ColumnHistogramStatsCacheLoader.java b/...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
@@ -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);
     }