[BugFix] Fix the mem_tracker use-after-free of UDF (backport StarRock…

…s#41710) (StarRocks#41727) Co-authored-by: trueeyu <lxhhust350@qq.com>
Seaven · Feb 27, 2024 · 41f0d7d · 41f0d7d
1 parent 4d1289c
commit 41f0d7d
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/be/src/udf/java/utils.cpp b/be/src/udf/java/utils.cpp
@@ -19,10 +19,14 @@ PromiseStatusPtr call_function_in_pthread(RuntimeState* state, const std::functi
     PromiseStatusPtr ms = std::make_unique<PromiseStatus>();
     if (bthread_self()) {
         state->exec_env()->udf_call_pool()->offer([promise = ms.get(), state, func]() {
-            MemTracker* prev_tracker = tls_thread_status.set_mem_tracker(state->instance_mem_tracker());
-            SCOPED_SET_TRACE_INFO({}, state->query_id(), state->fragment_instance_id());
-            DeferOp op([&] { tls_thread_status.set_mem_tracker(prev_tracker); });
-            promise->set_value(func());
+            Status st;
+            {
+                MemTracker* prev_tracker = tls_thread_status.set_mem_tracker(state->instance_mem_tracker());
+                SCOPED_SET_TRACE_INFO({}, state->query_id(), state->fragment_instance_id());
+                DeferOp op([&] { tls_thread_status.set_mem_tracker(prev_tracker); });
+                st = func();
+            }
+            promise->set_value(st);
         });
     } else {
         ms->set_value(func());