Skip to content

Commit 9559176

Browse files
mergify[bot]satanson
mergify[bot]
and
satanson
authored
[BugFix] Fix memory-leak introduced by udaf cache (backport #74025) (backport #74026) (#74251)
Signed-off-by: satanson <ranpanf@gmail.com> Co-authored-by: satanson <ranpanf@gmail.com>
1 parent dc09c5c commit 9559176

9 files changed

Lines changed: 51 additions & 94 deletions

‎be/src/exec/aggregator.cpp‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -285,8 +285,6 @@ Status Aggregator::open(RuntimeState* state) {
285285
auto& opts = state->query_options();
286286
bool enable_cache = opts.__isset.enable_cache_udaf && opts.enable_cache_udaf;
287287
auto promise_st = call_function_in_pthread(state, [this, enable_cache]() {
288-
std::vector<int> attached_udaf_idx;
289-
attached_udaf_idx.reserve(_agg_fn_ctxs.size());
290288
for (int i = 0; i < _agg_fn_ctxs.size(); ++i) {
291289
if (_fns[i].binary_type == TFunctionBinaryType::SRJAR) {
292290
const auto& fn = _fns[i];
@@ -307,13 +305,7 @@ Status Aggregator::open(RuntimeState* state) {
307305
COUNTER_UPDATE(_agg_stat->udaf_cache_populate_count, 1);
308306
}
309307
}
310-
if (!st.ok()) {
311-
for (int idx : attached_udaf_idx) {
312-
destroy_java_udaf_context(_agg_fn_ctxs[idx]);
313-
}
314-
return st;
315-
}
316-
attached_udaf_idx.emplace_back(i);
308+
RETURN_IF_ERROR(st);
317309
}
318310
}
319311
return Status::OK();

‎be/src/exprs/agg/java_udaf_function.cpp‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -122,9 +122,11 @@ static StatusOr<std::shared_ptr<JavaUDAFSharedContext>> build_udaf_shared_contex
122122
}
123123

124124
// Build a per-aggregator JavaUDAFUniqueContext on top of a (possibly cached) JavaUDAFSharedContext.
125-
static Status build_udaf_unique_context(std::shared_ptr<JavaUDAFSharedContext> udaf_ctx, FunctionContext* context) {
126-
auto agg_ctx = std::make_unique<JavaUDAFUniqueContext>();
127-
agg_ctx->ctx = std::move(udaf_ctx);
125+
// The context is already pre-allocated in FunctionContext::_jvm_udaf_ctxs; we populate it here.
126+
static Status build_udaf_unique_context(std::shared_ptr<JavaUDAFSharedContext> udaf_shared_ctx,
127+
FunctionContext* context) {
128+
auto* agg_ctx = context->udaf_ctxs();
129+
agg_ctx->ctx = std::move(udaf_shared_ctx);
128130

129131
// Create a per-aggregator UDAF object instance
130132
ASSIGN_OR_RETURN(agg_ctx->handle, agg_ctx->ctx->udaf_class.newInstance());
@@ -149,8 +151,7 @@ static Status build_udaf_unique_context(std::shared_ptr<JavaUDAFSharedContext> u
149151
JavaGlobalRef(env->NewGlobalRef(agg_ctx->ctx->states_add_method.handle())),
150152
JavaGlobalRef(env->NewGlobalRef(agg_ctx->ctx->states_remove_method.handle())),
151153
JavaGlobalRef(env->NewGlobalRef(agg_ctx->ctx->states_clear_method.handle())));
152-
agg_ctx->_func = std::make_unique<UDAFFunction>(agg_ctx->handle.handle(), context, agg_ctx.get());
153-
attach_java_udaf_context(context, std::move(agg_ctx));
154+
agg_ctx->_func = std::make_unique<UDAFFunction>(agg_ctx->handle.handle(), context, agg_ctx);
154155
return Status::OK();
155156
}
156157

‎be/src/exprs/agg/java_udaf_function.h‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
5757
input_column = down_cast<const BinaryColumn*>(column);
5858
}
5959
Slice slice = input_column->get_slice(row_num);
60-
auto* udaf_ctx = get_java_udaf_context(ctx);
60+
auto* udaf_ctx = ctx->udaf_ctxs();
6161

6262
if (udaf_ctx->buffer->capacity() < slice.get_size()) {
6363
udaf_ctx->buffer_data.resize(slice.get_size());
@@ -83,7 +83,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
8383
}
8484

8585
size_t old_size = column->get_bytes().size();
86-
auto* udaf_ctx = get_java_udaf_context(ctx);
86+
auto* udaf_ctx = ctx->udaf_ctxs();
8787
int serialize_size = udaf_ctx->_func->serialize_size(this->data(state).handle);
8888
if (udaf_ctx->buffer->capacity() < serialize_size) {
8989
udaf_ctx->buffer_data.resize(serialize_size);
@@ -101,7 +101,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
101101

102102
void finalize_to_column([[maybe_unused]] FunctionContext* ctx, ConstAggDataPtr __restrict state,
103103
Column* to) const final {
104-
auto* udaf_ctx = get_java_udaf_context(ctx);
104+
auto* udaf_ctx = ctx->udaf_ctxs();
105105
DCHECK(udaf_ctx != nullptr);
106106
jvalue val = udaf_ctx->_func->finalize(this->data(state).handle);
107107
// STRUCT-bearing return types route through append_jvalue's STRUCT path with the
@@ -117,7 +117,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
117117
ColumnPtr* dst) const final {
118118
auto& helper = JVMFunctionHelper::getInstance();
119119
auto* env = helper.getEnv();
120-
auto* udf_ctxs = get_java_udaf_context(ctx);
120+
auto* udf_ctxs = ctx->udaf_ctxs();
121121
// 1 convert input as state
122122
// 1.1 create state list
123123
auto rets =
@@ -192,12 +192,12 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
192192
// jclass
193193
// newInstance -> handle
194194
void create(FunctionContext* ctx, AggDataPtr __restrict ptr) const override {
195-
new (ptr) State(get_java_udaf_context(ctx)->_func->create());
195+
new (ptr) State(ctx->udaf_ctxs()->_func->create());
196196
}
197197

198198
// Call Destroy method
199199
void destroy(FunctionContext* ctx, AggDataPtr __restrict ptr) const override {
200-
get_java_udaf_context(ctx)->_func->destroy(data(ptr).handle);
200+
ctx->udaf_ctxs()->_func->destroy(data(ptr).handle);
201201
data(ptr).~State();
202202
}
203203

@@ -222,7 +222,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
222222
void update_batch(FunctionContext* ctx, size_t batch_size, size_t state_offset, const Column** columns,
223223
AggDataPtr* states) const override {
224224
auto& helper = JVMFunctionHelper::getInstance();
225-
auto* udf_ctxs = get_java_udaf_context(ctx);
225+
auto* udf_ctxs = ctx->udaf_ctxs();
226226
std::vector<jobject> args;
227227
int num_cols = ctx->get_num_args();
228228
helper.getEnv()->PushLocalFrame(num_cols * 3 + 1);
@@ -244,7 +244,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
244244
void update_batch_selectively(FunctionContext* ctx, size_t batch_size, size_t state_offset, const Column** columns,
245245
AggDataPtr* states, const Filter& filter) const override {
246246
auto [env, helper] = JVMFunctionHelper::getInstanceWithEnv();
247-
auto* udf_ctxs = get_java_udaf_context(ctx);
247+
auto* udf_ctxs = ctx->udaf_ctxs();
248248
std::vector<jobject> args;
249249
int num_cols = ctx->get_num_args();
250250
helper.getEnv()->PushLocalFrame(num_cols * 3 + 1);
@@ -269,7 +269,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
269269
auto* env = helper.getEnv();
270270
std::vector<jobject> args;
271271
int num_cols = ctx->get_num_args();
272-
auto* udf_ctxs = get_java_udaf_context(ctx);
272+
auto* udf_ctxs = ctx->udaf_ctxs();
273273
env->PushLocalFrame(num_cols * 3 + 1);
274274
auto defer = DeferOp([env = env]() { env->PopLocalFrame(nullptr); });
275275
{
@@ -279,7 +279,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
279279
SET_FUNCTION_CONTEXT_ERR(st, ctx);
280280
RETURN_IF_UNLIKELY(!st.ok(), (void)0);
281281

282-
auto* stub = get_java_udaf_context(ctx)->update_batch_call_stub.get();
282+
auto* stub = ctx->udaf_ctxs()->update_batch_call_stub.get();
283283
auto state_handle = this->data(state).handle;
284284
helper.batch_update_single(stub, state_handle, args.data(), num_cols, batch_size);
285285
}
@@ -320,7 +320,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
320320
// batch merge
321321
auto& helper = JVMFunctionHelper::getInstance();
322322
auto* env = helper.getEnv();
323-
auto* udf_ctxs = get_java_udaf_context(ctx);
323+
auto* udf_ctxs = ctx->udaf_ctxs();
324324

325325
auto provider = [&]() {
326326
auto state_id_list = JavaDataTypeConverter::convert_to_states(ctx, states, state_offset, batch_size);
@@ -341,7 +341,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
341341
AggDataPtr* states, const Filter& filter) const override {
342342
// batch merge
343343
auto& helper = JVMFunctionHelper::getInstance();
344-
auto* udf_ctxs = get_java_udaf_context(ctx);
344+
auto* udf_ctxs = ctx->udaf_ctxs();
345345

346346
auto provider = [&]() {
347347
auto state_id_list = JavaDataTypeConverter::convert_to_states_with_filter(ctx, states, state_offset,
@@ -360,7 +360,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
360360
size_t size) const override {
361361
auto& helper = JVMFunctionHelper::getInstance();
362362
auto* env = helper.getEnv();
363-
auto* udf_ctxs = get_java_udaf_context(ctx);
363+
auto* udf_ctxs = ctx->udaf_ctxs();
364364
auto provider = [&]() {
365365
auto state_handle = reinterpret_cast<JavaUDAFState*>(state)->handle;
366366
auto res = helper.convert_handle_to_jobject(ctx, state_handle);
@@ -379,7 +379,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
379379
size_t state_offset, Column* to) const override {
380380
auto& helper = JVMFunctionHelper::getInstance();
381381
auto* env = helper.getEnv();
382-
auto* udf_ctxs = get_java_udaf_context(ctx);
382+
auto* udf_ctxs = ctx->udaf_ctxs();
383383

384384
const size_t origin_chunk_size = to->size();
385385
auto defer = DeferOp([&]() {
@@ -435,7 +435,7 @@ class JavaUDAFAggregateFunction : public AggregateFunction {
435435
size_t state_offset, Column* to) const override {
436436
auto& helper = JVMFunctionHelper::getInstance();
437437
auto* env = helper.getEnv();
438-
auto* udf_ctxs = get_java_udaf_context(ctx);
438+
auto* udf_ctxs = ctx->udaf_ctxs();
439439

440440
const size_t origin_chunk_size = to->size();
441441
auto defer = DeferOp([&]() {

‎be/src/exprs/agg/java_window_function.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,9 @@ static StatusOr<std::shared_ptr<JavaUDAFSharedContext>> build_window_shared_cont
7171
}
7272

7373
// Build a per-aggregator JavaUDAFUniqueContext for a window function on top of a shared context.
74+
// The context is already pre-allocated in FunctionContext::_jvm_udaf_ctxs; we populate it here.
7475
static Status build_window_unique_context(std::shared_ptr<JavaUDAFSharedContext> shared, FunctionContext* context) {
75-
auto udaf_ctx = std::make_unique<JavaUDAFUniqueContext>();
76+
auto* udaf_ctx = context->udaf_ctxs();
7677
udaf_ctx->ctx = std::move(shared);
7778

7879
ASSIGN_OR_RETURN(udaf_ctx->handle, udaf_ctx->ctx->udaf_class.newInstance());
@@ -87,8 +88,7 @@ static Status build_window_unique_context(std::shared_ptr<JavaUDAFSharedContext>
8788
JavaGlobalRef(env->NewGlobalRef(udaf_ctx->ctx->states_add_method.handle())),
8889
JavaGlobalRef(env->NewGlobalRef(udaf_ctx->ctx->states_remove_method.handle())),
8990
JavaGlobalRef(env->NewGlobalRef(udaf_ctx->ctx->states_clear_method.handle())));
90-
udaf_ctx->_func = std::make_unique<UDAFFunction>(udaf_ctx->handle.handle(), context, udaf_ctx.get());
91-
attach_java_udaf_context(context, std::move(udaf_ctx));
91+
udaf_ctx->_func = std::make_unique<UDAFFunction>(udaf_ctx->handle.handle(), context, udaf_ctx);
9292
return Status::OK();
9393
}
9494

‎be/src/exprs/agg/java_window_function.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ Status assign_jvalue(const TypeDescriptor& type_desc, bool is_box, Column* col,
3434
class JavaWindowFunction final : public JavaUDAFAggregateFunction {
3535
public:
3636
void reset(FunctionContext* ctx, const Columns& args, AggDataPtr __restrict state) const override {
37-
get_java_udaf_context(ctx)->_func->reset(data(state).handle);
37+
ctx->udaf_ctxs()->_func->reset(data(state).handle);
3838
}
3939

4040
std::string get_name() const override { return "java_window"; }
@@ -64,8 +64,8 @@ class JavaWindowFunction final : public JavaUDAFAggregateFunction {
6464
SET_FUNCTION_CONTEXT_ERR(st, ctx);
6565
RETURN_IF_UNLIKELY(!st.ok(), (void)0);
6666

67-
get_java_udaf_context(ctx)->_func->window_update_batch(data(state).handle, peer_group_start, peer_group_end,
68-
frame_start, frame_end, num_args, args.data());
67+
ctx->udaf_ctxs()->_func->window_update_batch(data(state).handle, peer_group_start, peer_group_end, frame_start,
68+
frame_end, num_args, args.data());
6969
// release input cols
7070
for (int i = 0; i < num_args; ++i) {
7171
env->DeleteLocalRef(args[i]);
@@ -75,7 +75,7 @@ class JavaWindowFunction final : public JavaUDAFAggregateFunction {
7575
void get_values(FunctionContext* ctx, ConstAggDataPtr __restrict state, Column* dst, size_t start,
7676
size_t end) const override {
7777
auto& helper = JVMFunctionHelper::getInstance();
78-
jvalue val = get_java_udaf_context(ctx)->_func->finalize(this->data(state).handle);
78+
jvalue val = ctx->udaf_ctxs()->_func->finalize(this->data(state).handle);
7979
// insert values to column
8080
JNIEnv* env = helper.getEnv();
8181
const auto& return_type = ctx->get_return_type();

‎be/src/exprs/function_context.cpp‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ FunctionContext* FunctionContext::create_context(RuntimeState* state, MemPool* p
3838
ctx->_mem_pool = pool;
3939
ctx->_return_type = return_type;
4040
ctx->_arg_types = arg_types;
41+
#if !defined(BUILD_FORMAT_LIB)
42+
ctx->_jvm_udaf_ctxs = std::make_unique<JavaUDAFUniqueContext>();
43+
#endif
4144
return ctx;
4245
}
4346

@@ -51,6 +54,9 @@ FunctionContext* FunctionContext::create_context(RuntimeState* state, MemPool* p
5154
ctx->_mem_pool = pool;
5255
ctx->_return_type = return_type;
5356
ctx->_arg_types = arg_types;
57+
#if !defined(BUILD_FORMAT_LIB)
58+
ctx->_jvm_udaf_ctxs = std::make_unique<JavaUDAFUniqueContext>();
59+
#endif
5460
ctx->_is_distinct = is_distinct;
5561
ctx->_is_asc_order = is_asc_order;
5662
ctx->_nulls_first = nulls_first;
@@ -131,11 +137,12 @@ void* FunctionContext::get_function_state(FunctionStateScope scope) const {
131137
}
132138

133139
void FunctionContext::release_mems() {
134-
auto* udaf_ctx = get_java_udaf_context(this);
135-
if (udaf_ctx != nullptr && udaf_ctx->states) {
140+
#if !defined(BUILD_FORMAT_LIB)
141+
if (_jvm_udaf_ctxs != nullptr && _jvm_udaf_ctxs->states) {
136142
auto env = JVMFunctionHelper::getInstance().getEnv();
137-
udaf_ctx->states->clear(this, env);
143+
_jvm_udaf_ctxs->states->clear(this, env);
138144
}
145+
#endif
139146
}
140147

141148
void FunctionContext::set_error(const char* error_msg, const bool is_udf) {

‎be/src/exprs/function_context.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ class RuntimeState;
3535
class Column;
3636
class Slice;
3737
struct NgramBloomFilterState;
38+
struct JavaUDAFUniqueContext;
3839

3940
class FunctionContext {
4041
public:
@@ -163,6 +164,8 @@ class FunctionContext {
163164
bool has_error() const;
164165
const char* error_msg() const;
165166

167+
JavaUDAFUniqueContext* udaf_ctxs() { return _jvm_udaf_ctxs.get(); }
168+
166169
void release_mems();
167170

168171
ssize_t get_group_concat_max_len() { return group_concat_max_len; }
@@ -214,6 +217,9 @@ class FunctionContext {
214217
// it will point to the internal _mem_usage
215218
int64_t* _mem_usage_counter = &_mem_usage;
216219

220+
// UDAF Context
221+
std::unique_ptr<JavaUDAFUniqueContext> _jvm_udaf_ctxs;
222+
217223
std::vector<bool> _is_asc_order;
218224
std::vector<bool> _nulls_first;
219225
bool _is_distinct = false;

‎be/src/udf/java/java_udf.cpp‎

Lines changed: 5 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -120,49 +120,6 @@ static JNINativeMethod java_native_methods[] = {
120120
};
121121
#pragma GCC diagnostic pop
122122

123-
JavaUDAFUniqueContext* get_java_udaf_context(FunctionContext* ctx) {
124-
if (ctx == nullptr) {
125-
return nullptr;
126-
}
127-
return reinterpret_cast<JavaUDAFUniqueContext*>(ctx->get_function_state(FunctionContext::THREAD_LOCAL));
128-
}
129-
130-
void attach_java_udaf_context(FunctionContext* ctx, std::unique_ptr<JavaUDAFUniqueContext> udaf_ctx) {
131-
DCHECK(ctx != nullptr);
132-
DCHECK(udaf_ctx != nullptr);
133-
auto* old_ctx = get_java_udaf_context(ctx);
134-
DCHECK(old_ctx == nullptr) << "duplicate Java UDAF context attach";
135-
if (old_ctx != nullptr) {
136-
delete old_ctx;
137-
}
138-
ctx->set_function_state(FunctionContext::THREAD_LOCAL, udaf_ctx.release());
139-
}
140-
141-
void clear_java_udaf_states(FunctionContext* ctx) {
142-
auto* udaf_ctx = get_java_udaf_context(ctx);
143-
if (udaf_ctx == nullptr || udaf_ctx->states == nullptr) {
144-
return;
145-
}
146-
147-
auto env = JVMFunctionHelper::getInstance().getEnv();
148-
udaf_ctx->states->clear(ctx, env);
149-
}
150-
151-
void destroy_java_udaf_context(FunctionContext* ctx) {
152-
if (ctx == nullptr) {
153-
return;
154-
}
155-
156-
auto* udaf_ctx = get_java_udaf_context(ctx);
157-
if (udaf_ctx == nullptr) {
158-
return;
159-
}
160-
161-
clear_java_udaf_states(ctx);
162-
ctx->set_function_state(FunctionContext::THREAD_LOCAL, nullptr);
163-
delete udaf_ctx;
164-
}
165-
166123
StatusOr<jobject> MapMeta::newLocalInstance(jobject keys, jobject values) const {
167124
JNIEnv* env = getJNIEnv();
168125
auto res = env->NewObject(immutable_map_class->clazz(), immutable_map_constructor, keys, values);
@@ -540,8 +497,8 @@ void JVMFunctionHelper::batch_update(FunctionContext* ctx, jobject udaf, jobject
540497
int cols) {
541498
jobjectArray input_arr = _build_object_array(_object_array_class, input, cols);
542499
LOCAL_REF_GUARD(input_arr);
543-
_env->CallStaticVoidMethod(_udf_helper_class, _batch_update, udaf, update,
544-
get_java_udaf_context(ctx)->states->handle(), states, input_arr);
500+
_env->CallStaticVoidMethod(_udf_helper_class, _batch_update, udaf, update, ctx->udaf_ctxs()->states->handle(),
501+
states, input_arr);
545502
CHECK_UDF_CALL_EXCEPTION(_env, ctx);
546503
}
547504

@@ -558,7 +515,7 @@ void JVMFunctionHelper::batch_update_if_not_null(FunctionContext* ctx, jobject u
558515
jobjectArray input_arr = _build_object_array(_object_array_class, input, cols);
559516
LOCAL_REF_GUARD(input_arr);
560517
_env->CallStaticVoidMethod(_udf_helper_class, _batch_update_if_not_null, udaf, update,
561-
get_java_udaf_context(ctx)->states->handle(), states, input_arr);
518+
ctx->udaf_ctxs()->states->handle(), states, input_arr);
562519
CHECK_UDF_CALL_EXCEPTION(_env, ctx);
563520
}
564521

@@ -620,12 +577,12 @@ void JVMFunctionHelper::get_result_from_boxed_array(FunctionContext* ctx, int ty
620577

621578
// convert UDAF ctx to jobject
622579
jobject JVMFunctionHelper::convert_handle_to_jobject(FunctionContext* ctx, int state) {
623-
auto* states = get_java_udaf_context(ctx)->states.get();
580+
auto* states = ctx->udaf_ctxs()->states.get();
624581
return states->get_state(ctx, _env, state);
625582
}
626583

627584
jobject JVMFunctionHelper::convert_handles_to_jobjects(FunctionContext* ctx, jobject state_ids) {
628-
auto* states = get_java_udaf_context(ctx)->states.get();
585+
auto* states = ctx->udaf_ctxs()->states.get();
629586
return states->get_state(ctx, _env, state_ids);
630587
}
631588

0 commit comments

Comments
 (0)