(auron) branch master updated: [AURON #2119] Add `isTaskRunning` method to check task status in SparkAuronAdaptor (#2120)

[csy]

This is an automated email from the ASF dual-hosted git repository.

csy pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git


The following commit(s) were added to refs/heads/master by this push:
     new cb81ca14 [AURON #2119] Add `isTaskRunning` method to check task status 
in SparkAuronAdaptor (#2120)
cb81ca14 is described below

commit cb81ca14007fcdaf27bfcc558bdce9f980c4933e
Author: cxzl25 <3898450+cxz...@users.noreply.github.com>
AuthorDate: Mon Mar 30 10:53:40 2026 +0800

    [AURON #2119] Add `isTaskRunning` method to check task status in 
SparkAuronAdaptor (#2120)
    
    # Which issue does this PR close?
    
    Closes #2119
    
    # Rationale for this change
    
    When Spark initiates a task kill (e.g. stage cancellation, speculative
    execution kill, or user-triggered job cancellation), the Rust native
    engine always receives true from the JNI callback
    JniBridge.isTaskRunning(), making it unable to detect that the task has
    already been killed.
    
    # What changes are included in this PR?
    
    # Are there any user-facing changes?
    
    # How was this patch tested?
---
 .../src/main/java/org/apache/auron/jni/SparkAuronAdaptor.java    | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git 
a/spark-extension/src/main/java/org/apache/auron/jni/SparkAuronAdaptor.java 
b/spark-extension/src/main/java/org/apache/auron/jni/SparkAuronAdaptor.java
index 0d90957b..c58b8c69 100644
--- a/spark-extension/src/main/java/org/apache/auron/jni/SparkAuronAdaptor.java
+++ b/spark-extension/src/main/java/org/apache/auron/jni/SparkAuronAdaptor.java
@@ -52,6 +52,15 @@ public class SparkAuronAdaptor extends AuronAdaptor {
         }
     }
 
+    @Override
+    public boolean isTaskRunning() {
+        TaskContext tc = TaskContext$.MODULE$.get();
+        if (tc == null) { // driver is always running
+            return true;
+        }
+        return !tc.isCompleted() && !tc.isInterrupted();
+    }
+
     @Override
     public Object getThreadContext() {
         return TaskContext$.MODULE$.get();