From 696b750e7cca609734192a2a39342bd99c753d83 Mon Sep 17 00:00:00 2001 From: taiyang-li Date: Thu, 2 Jul 2026 21:24:37 +0800 Subject: [PATCH] [GLUTEN][CORE] Add TaskContext JNI callback for reading Spark task attempt id from native Instead of extending every JNI entry point (createRuntime / MemoryManager create/hold/release) to plumb the Spark task attempt id from Java down to C++ as an extra parameter, expose a small callback surface that native code uses on demand: - Java side: org.apache.gluten.task.TaskContextJniWrapper#currentTaskAttemptId() reads TaskContext.get().taskAttemptId() on the current thread and returns -1 when there is no task context. - C++ side: gluten::getCurrentSparkTaskAttemptId() attaches the current thread to the JVM as a daemon on demand and calls back into the Java helper via JNI. The class ref and method id are cached in function-local statics on first use. Because Spark's TaskContext is a per-thread ThreadLocal, this returns a meaningful value whenever the native call is running on an executor task thread (or any thread inheriting that ThreadLocal), which is exactly when backends need it. No signature change to Runtime / MemoryManager / RuntimeJniWrapper / NativeMemoryManagerJniWrapper. No behavior change for existing backends (Velox, ClickHouse) that do not query the task attempt id from native. Co-Authored-By: Aime Change-Id: I3185249796b0c396813dc39f54bd8e8b8589ca2a --- cpp/core/CMakeLists.txt | 3 +- cpp/core/jni/JniCommon.h | 4 ++ cpp/core/jni/TaskContextJniWrapper.cc | 38 +++++++++++++++++ cpp/core/jni/TaskContextJniWrapper.h | 33 +++++++++++++++ .../gluten/task/TaskContextJniWrapper.java | 42 +++++++++++++++++++ 5 files changed, 119 insertions(+), 1 deletion(-) create mode 100755 cpp/core/jni/TaskContextJniWrapper.cc create mode 100755 cpp/core/jni/TaskContextJniWrapper.h create mode 100755 gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java diff --git a/cpp/core/CMakeLists.txt b/cpp/core/CMakeLists.txt index 9530d603f5b..f3860166a49 100644 --- a/cpp/core/CMakeLists.txt +++ b/cpp/core/CMakeLists.txt @@ -150,7 +150,8 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS utils/StringUtil.cc utils/ObjectStore.cc jni/JniError.cc - jni/JniCommon.cc) + jni/JniCommon.cc + jni/TaskContextJniWrapper.cc) file(MAKE_DIRECTORY ${root_directory}/releases) add_library(gluten SHARED ${SPARK_COLUMNAR_PLUGIN_SRCS}) diff --git a/cpp/core/jni/JniCommon.h b/cpp/core/jni/JniCommon.h index 43ddb2226a0..bf5a6a746b7 100644 --- a/cpp/core/jni/JniCommon.h +++ b/cpp/core/jni/JniCommon.h @@ -163,6 +163,10 @@ class JniCommonState { jmethodID runtimeAwareCtxHandle(); + JavaVM* getJavaVM() const { + return vm_; + } + private: void initialize(JNIEnv* env); diff --git a/cpp/core/jni/TaskContextJniWrapper.cc b/cpp/core/jni/TaskContextJniWrapper.cc new file mode 100755 index 00000000000..4349f1cbf13 --- /dev/null +++ b/cpp/core/jni/TaskContextJniWrapper.cc @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "jni/TaskContextJniWrapper.h" + +#include "jni/JniCommon.h" + +namespace gluten { + +int64_t getCurrentSparkTaskAttemptId() { + JavaVM* vm = getJniCommonState()->getJavaVM(); + JNIEnv* env = nullptr; + attachCurrentThreadAsDaemonOrThrow(vm, &env); + + static jclass taskContextJniWrapperClass = + createGlobalClassReferenceOrError(env, "Lorg/apache/gluten/task/TaskContextJniWrapper;"); + static jmethodID currentTaskAttemptIdMethod = + getStaticMethodIdOrError(env, taskContextJniWrapperClass, "currentTaskAttemptId", "()J"); + + jlong id = env->CallStaticLongMethod(taskContextJniWrapperClass, currentTaskAttemptIdMethod); + checkException(env); + return static_cast(id); +} + +} // namespace gluten diff --git a/cpp/core/jni/TaskContextJniWrapper.h b/cpp/core/jni/TaskContextJniWrapper.h new file mode 100755 index 00000000000..87cff54c55e --- /dev/null +++ b/cpp/core/jni/TaskContextJniWrapper.h @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace gluten { + +/** + * Returns the current thread's Spark task attempt id by calling back into + * org.apache.gluten.task.TaskContextJniWrapper#currentTaskAttemptId via JNI. + * + * Returns -1 when there is no Spark TaskContext on the calling thread (driver, + * standalone native worker, etc.). Safe to call from any thread; the helper + * attaches the current thread to the JVM as a daemon on demand. + */ +int64_t getCurrentSparkTaskAttemptId(); + +} // namespace gluten diff --git a/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java b/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java new file mode 100755 index 00000000000..794ee9cb65b --- /dev/null +++ b/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.task; + +import org.apache.spark.TaskContext; + +/** + * Callback surface that native code uses to obtain the current thread's Spark task attempt id via + * JNI, without having to plumb it explicitly through every JNI entry point. + * + *

Native side reads the value from {@link TaskContext#get()} at the moment it is needed. Since + * {@link TaskContext} is a per-thread ThreadLocal, this only returns a meaningful value when the + * native call is running on the executor task thread (or on a JVM thread inheriting the same + * ThreadLocal). It returns {@code -1L} on the driver or on a native worker thread that has no + * associated Spark task. + */ +public final class TaskContextJniWrapper { + private TaskContextJniWrapper() {} + + /** + * Returns the current thread's Spark task attempt id, or {@code -1L} when there is no task + * context on the calling thread. + */ + public static long currentTaskAttemptId() { + TaskContext tc = TaskContext.get(); + return tc == null ? -1L : tc.taskAttemptId(); + } +}