(airflow) branch main updated: modified cond statements when heatbeat returned cancellation (#68535)

[jasonliu]

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

jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new f533c0cfdb4 modified cond statements when heatbeat returned 
cancellation (#68535)
f533c0cfdb4 is described below

commit f533c0cfdb42dd3e0f9259345d29ddf0fa82668b
Author: jj.lee <63435794+jx2...@users.noreply.github.com>
AuthorDate: Sun Jun 28 18:51:51 2026 +0900

    modified cond statements when heatbeat returned cancellation (#68535)
---
 go-sdk/pkg/worker/runner.go      |  2 +-
 go-sdk/pkg/worker/runner_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/go-sdk/pkg/worker/runner.go b/go-sdk/pkg/worker/runner.go
index 4065e30ad70..d30c2351bbe 100644
--- a/go-sdk/pkg/worker/runner.go
+++ b/go-sdk/pkg/worker/runner.go
@@ -327,7 +327,7 @@ func (w *worker) ExecuteTaskWorkload(ctx context.Context, 
workload api.ExecuteTa
        var finalState api.TerminalTIState
        body := &api.TIUpdateStatePayload{}
 
-       if taskContext.Err() == ErrTaskCancelledAfterFailedHeartbeat {
+       if errors.Is(context.Cause(taskContext), 
ErrTaskCancelledAfterFailedHeartbeat) {
                // We've already logged when we failed to heartbeat, don't do 
it again
                finalState = api.TerminalTIStateFailed
                body.FromTITerminalStatePayload(api.TITerminalStatePayload{
diff --git a/go-sdk/pkg/worker/runner_test.go b/go-sdk/pkg/worker/runner_test.go
index 72673c00df3..01ab4ba62b2 100644
--- a/go-sdk/pkg/worker/runner_test.go
+++ b/go-sdk/pkg/worker/runner_test.go
@@ -21,6 +21,7 @@ import (
        "context"
        "fmt"
        "log/slog"
+       "net/http"
        "sync/atomic"
        "testing"
        "time"
@@ -31,6 +32,7 @@ import (
        "github.com/spf13/viper"
        "github.com/stretchr/testify/mock"
        "github.com/stretchr/testify/suite"
+       "resty.dev/v3"
 
        "github.com/apache/airflow/go-sdk/bundle/bundlev1"
        "github.com/apache/airflow/go-sdk/pkg/api"
@@ -218,6 +220,38 @@ func (s *WorkerSuite) TestTaskHeartbeatsWhileRunning() {
                True(count <= 11 && count >= 9, fmt.Sprintf("Call count of %d 
was not within the margin of error of 10+/-1", count))
 }
 
+func (s *WorkerSuite) TestTaskHeartbeatConflictStopsTask() {
+       id := uuid.New().String()
+       testWorkload := newTestWorkLoad(id, id[:8])
+
+       s.registry.AddDag(testWorkload.TI.DagId).
+               AddTaskWithName(testWorkload.TI.TaskId, func(ctx 
context.Context) error {
+                       select {
+                       case <-ctx.Done():
+                               return nil
+                       case <-time.After(2 * time.Second):
+                               return fmt.Errorf("task context was not 
cancelled")
+                       }
+               })
+
+       s.ExpectTaskRun(id)
+       s.ExpectTaskState(id, api.TerminalTIStateFailed)
+       s.ti.EXPECT().
+               Heartbeat(mock.Anything, uuid.MustParse(id), mock.Anything).
+               Return(&api.GeneralHTTPError{
+                       Response: &resty.Response{
+                               RawResponse: &http.Response{
+                                       Status:     "409 Conflict",
+                                       StatusCode: http.StatusConflict,
+                               },
+                       },
+               })
+       s.client.EXPECT().TaskInstances().Return(s.ti)
+
+       err := s.worker.ExecuteTaskWorkload(context.Background(), testWorkload)
+       s.NoError(err)
+}
+
 func (s *WorkerSuite) TestTaskHeartbeatErrorStopsTaskAndLogs() {
        s.T().Skip("TODO: Not implemented yet")
 }