From 164761bb5315bce344e16b30192bcb942a30a659 Mon Sep 17 00:00:00 2001 From: trueeyu Date: Mon, 26 Feb 2024 18:34:30 +0800 Subject: [PATCH] Fix the mem_tracker use-after-free of UDF Signed-off-by: trueeyu --- be/src/udf/java/utils.cpp | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/be/src/udf/java/utils.cpp b/be/src/udf/java/utils.cpp index 705c64cf3a917..05b14ca63268c 100644 --- a/be/src/udf/java/utils.cpp +++ b/be/src/udf/java/utils.cpp @@ -32,10 +32,14 @@ PromiseStatusPtr call_function_in_pthread(RuntimeState* state, const std::functi PromiseStatusPtr ms = std::make_unique(); 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());