[GLUTEN][CORE] Add TaskContext JNI callback for reading Spark task attempt id from native (#12435)

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.


Change-Id: I3185249796b0c396813dc39f54bd8e8b8589ca2a

Co-authored-by: Aime <aime@bytedance.com>

Author: taiyang-li

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})

cpp/core/jni/JniCommon.h

@@ -163,6 +163,10 @@ class JniCommonState {
 
   jmethodID runtimeAwareCtxHandle();
 
+  JavaVM* getJavaVM() const {
+    return vm_;
+  }
+
  private:
   void initialize(JNIEnv* env);
 

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<int64_t>(id);
+}
+
+} // namespace gluten

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 <cstdint>
+
+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

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