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(); + } +}