This is an automated email from the ASF dual-hosted git repository. kxiao pushed a commit to branch branch-2.0 in repository https://gitbox.apache.org/repos/asf/doris.git
commit dbff64ccbbb02885f895d63cdef7670d50e8320a Author: zhangstar333 <[email protected]> AuthorDate: Mon Oct 9 15:10:56 2023 +0800 [Bug](java-udf) fix java-udf memory leak (#25151) --- be/src/vec/functions/function_java_udf.cpp | 56 ++++++++++++++++-------------- be/src/vec/functions/function_java_udf.h | 4 --- 2 files changed, 29 insertions(+), 31 deletions(-) diff --git a/be/src/vec/functions/function_java_udf.cpp b/be/src/vec/functions/function_java_udf.cpp index fa9c8517c8e..331bb1bcac7 100644 --- a/be/src/vec/functions/function_java_udf.cpp +++ b/be/src/vec/functions/function_java_udf.cpp @@ -85,39 +85,41 @@ Status JavaFunctionCall::open(FunctionContext* context, FunctionContext::Functio RETURN_ERROR_IF_EXC(env); context->set_function_state(FunctionContext::FRAGMENT_LOCAL, jni_env); } - JniEnv* jni_env = - reinterpret_cast<JniEnv*>(context->get_function_state(FunctionContext::FRAGMENT_LOCAL)); - std::shared_ptr<JniContext> jni_ctx = std::make_shared<JniContext>( - _argument_types.size(), jni_env->executor_cl, jni_env->executor_close_id); - context->set_function_state(FunctionContext::THREAD_LOCAL, jni_ctx); - // Add a scoped cleanup jni reference object. This cleans up local refs made below. - JniLocalFrame jni_frame; - { - std::string local_location; - auto function_cache = UserFunctionCache::instance(); - RETURN_IF_ERROR(function_cache->get_jarpath(fn_.id, fn_.hdfs_location, fn_.checksum, - &local_location)); - TJavaUdfExecutorCtorParams ctor_params; - ctor_params.__set_fn(fn_); - ctor_params.__set_location(local_location); - jbyteArray ctor_params_bytes; + if (scope == FunctionContext::FunctionStateScope::THREAD_LOCAL) { + JniEnv* jni_env = reinterpret_cast<JniEnv*>( + context->get_function_state(FunctionContext::FRAGMENT_LOCAL)); + std::shared_ptr<JniContext> jni_ctx = std::make_shared<JniContext>( + _argument_types.size(), jni_env->executor_cl, jni_env->executor_close_id); + context->set_function_state(FunctionContext::THREAD_LOCAL, jni_ctx); - // Pushed frame will be popped when jni_frame goes out-of-scope. - RETURN_IF_ERROR(jni_frame.push(env)); + // Add a scoped cleanup jni reference object. This cleans up local refs made below. + JniLocalFrame jni_frame; + { + std::string local_location; + auto function_cache = UserFunctionCache::instance(); + RETURN_IF_ERROR(function_cache->get_jarpath(fn_.id, fn_.hdfs_location, fn_.checksum, + &local_location)); + TJavaUdfExecutorCtorParams ctor_params; + ctor_params.__set_fn(fn_); + ctor_params.__set_location(local_location); + jbyteArray ctor_params_bytes; - RETURN_IF_ERROR(SerializeThriftMsg(env, &ctor_params, &ctor_params_bytes)); + // Pushed frame will be popped when jni_frame goes out-of-scope. + RETURN_IF_ERROR(jni_frame.push(env)); - jni_ctx->executor = - env->NewObject(jni_env->executor_cl, jni_env->executor_ctor_id, ctor_params_bytes); + RETURN_IF_ERROR(SerializeThriftMsg(env, &ctor_params, &ctor_params_bytes)); - jbyte* pBytes = env->GetByteArrayElements(ctor_params_bytes, nullptr); - env->ReleaseByteArrayElements(ctor_params_bytes, pBytes, JNI_ABORT); - env->DeleteLocalRef(ctor_params_bytes); - } - RETURN_ERROR_IF_EXC(env); - RETURN_IF_ERROR(JniUtil::LocalToGlobalRef(env, jni_ctx->executor, &jni_ctx->executor)); + jni_ctx->executor = env->NewObject(jni_env->executor_cl, jni_env->executor_ctor_id, + ctor_params_bytes); + jbyte* pBytes = env->GetByteArrayElements(ctor_params_bytes, nullptr); + env->ReleaseByteArrayElements(ctor_params_bytes, pBytes, JNI_ABORT); + env->DeleteLocalRef(ctor_params_bytes); + } + RETURN_ERROR_IF_EXC(env); + RETURN_IF_ERROR(JniUtil::LocalToGlobalRef(env, jni_ctx->executor, &jni_ctx->executor)); + } return Status::OK(); } diff --git a/be/src/vec/functions/function_java_udf.h b/be/src/vec/functions/function_java_udf.h index 875e3426df1..d05e2e581f1 100644 --- a/be/src/vec/functions/function_java_udf.h +++ b/be/src/vec/functions/function_java_udf.h @@ -165,10 +165,6 @@ private: env->DeleteGlobalRef(executor); is_closed = true; } - - /// These functions are cross-compiled to IR and used by codegen. - static void SetInputNullsBufferElement(JniContext* jni_ctx, int index, uint8_t value); - static uint8_t* GetInputValuesBufferAtOffset(JniContext* jni_ctx, int offset); }; }; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
