[BugFix] remove threadLocal after operation (#34758)

Signed-off-by: packy92 <wangchao@starrocks.com>
StarRocks · Nov 15, 2023 · 32ff703 · 32ff703
1 parent b5848a2
commit 32ff703
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 1 deletion.
diff --git a/fe/fe-core/src/main/java/com/starrocks/connector/ConnectorColumnStatsCacheLoader.java b/fe/fe-core/src/main/java/com/starrocks/connector/ConnectorColumnStatsCacheLoader.java
@@ -66,6 +66,8 @@ public class ConnectorColumnStatsCacheLoader implements
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
@@ -103,6 +105,8 @@ public class ConnectorColumnStatsCacheLoader implements
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }

diff --git a/fe/fe-core/src/main/java/com/starrocks/http/BaseAction.java b/fe/fe-core/src/main/java/com/starrocks/http/BaseAction.java
@@ -41,6 +41,7 @@
 import com.starrocks.privilege.AuthorizationMgr;
 import com.starrocks.privilege.PrivilegeBuiltinConstants;
 import com.starrocks.privilege.PrivilegeException;
+import com.starrocks.qe.ConnectContext;
 import com.starrocks.server.GlobalStateMgr;
 import com.starrocks.sql.ast.UserIdentity;
 import io.netty.buffer.ByteBuf;
@@ -111,6 +112,8 @@ public void handleRequest(BaseRequest request) throws Exception {
             } else {
                 writeResponse(request, response, HttpResponseStatus.NOT_FOUND);
             }
+        } finally {
+            ConnectContext.remove();
         }
     }
 

diff --git a/fe/fe-core/src/main/java/com/starrocks/http/rest/ExecuteSqlAction.java b/fe/fe-core/src/main/java/com/starrocks/http/rest/ExecuteSqlAction.java
@@ -49,6 +49,7 @@
 import com.starrocks.http.HttpConnectContext;
 import com.starrocks.http.HttpConnectProcessor;
 import com.starrocks.http.IllegalArgException;
+import com.starrocks.qe.ConnectContext;
 import com.starrocks.qe.ConnectScheduler;
 import com.starrocks.qe.OriginStatement;
 import com.starrocks.qe.QueryState;
@@ -145,6 +146,8 @@ private void realWork(BaseRequest request, BaseResponse response) {
             } catch (Exception e) {
                 // just for safe. most Exception is handled in execute(), and set error code in context
                 throw new StarRocksHttpException(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage());
+            } finally {
+                ConnectContext.remove();
             }
 
             // finalize just send 200 for kill, and throw StarRocksHttpException if context's error is set

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
@@ -1138,7 +1138,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
@@ -69,6 +69,8 @@ public class ColumnBasicStatsCacheLoader implements AsyncCacheLoader<ColumnStats
                 throw e;
             } catch (Exception e) {
                 throw new CompletionException(e);
+            } finally {
+                ConnectContext.remove();
             }
         }, executor);
     }
@@ -106,6 +108,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);
     }