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 5b4f23b1c8f Go SDK: fail tasks that return non-JSON values instead of
hanging (#75601)
5b4f23b1c8f is described below
commit 5b4f23b1c8f78834bbff2b175e9ca031462d9744
Author: Y-C <[email protected]>
AuthorDate: Mon Nov 2 09:40:29 2026 +0800
Go SDK: fail tasks that return non-JSON values instead of hanging (#75601)
* Go SDK: fail tasks that return non-JSON values instead of hanging
The supervisor drops a SetXCom request whose value is not JSON without
replying, so a Go task that returned a time.Time, uuid.UUID or []byte
waited for that reply forever, or until execution_timeout when one is
set. The task state store already rejects such values before sending
them; XCom pushes need the same check, except that null stays allowed
because the Execution API stores it.
Putting every return value through that check also exposed a panic in
it: msgpack's default map decoder crashes on a key that decodes to
bytes, such as a uuid.UUID.
* Go SDK: test that PushXCom still pushes a nil value
A task that returns a nil slice, or nil as an any, pushes null as its
return value. Nothing failed when PushXCom rejected nil as well, so
moving the task state store's nil rule back into the shared check would
have failed those tasks with the whole suite still passing.
* Apply jason's suggestion
Co-authored-by: Jason(Zhe-You) Liu
<[email protected]>
---------
Co-authored-by: Eason09053360
<[email protected]>
Co-authored-by: Jason(Zhe-You) Liu
<[email protected]>
---
go-sdk/pkg/execution/client.go | 76 ++++++++++++++++++++++++++-----------
go-sdk/pkg/execution/client_test.go | 76 +++++++++++++++++++++++++++++++++++++
go-sdk/sdk/sdk.go | 4 ++
3 files changed, 134 insertions(+), 22 deletions(-)
diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go
index 69a63b0eee3..ba78706b97b 100644
--- a/go-sdk/pkg/execution/client.go
+++ b/go-sdk/pkg/execution/client.go
@@ -316,16 +316,22 @@ func (c *CoordinatorClient) PushXCom(
key string,
value any,
) error {
+ // Unlike a task state value, nil is allowed: it is pushed as JSON null.
+ encoded, _, err := validateJSONRepresentable(value)
+ if err != nil {
+ return fmt.Errorf("cannot push XCom %q: %w", key, err)
+ }
+
msg := genmodels.SetXCom{
Key: key,
- Value: value,
+ Value: encoded,
DagID: ti.DagID,
TaskID: ti.TaskID,
RunID: ti.RunID,
}
msg.MapIndex = omittedMapIndex(ti.MapIndex)
- _, err := c.comm.Communicate(ctx, msg)
+ _, err = c.comm.Communicate(ctx, msg)
return err
}
@@ -469,9 +475,15 @@ func (s taskStateStore) Set(
// The caller's value is checked first so a misconfigured deployment
cannot
// mask a programming error in the task.
- if err := validateJSONRepresentable(value); err != nil {
+ encoded, isNull, err := validateJSONRepresentable(value)
+ if err != nil {
return fmt.Errorf("cannot set task state key %q: %w", key, err)
}
+ // Checked after encoding: a typed nil such as a nil *string is not ==
nil,
+ // yet it still encodes to null, which the Execution API rejects.
+ if isNull {
+ return fmt.Errorf("cannot set task state key %q: value must not
be nil", key)
+ }
expiry, err := resolveExpiry(options.Retention, time.Now())
if err != nil {
@@ -481,7 +493,7 @@ func (s taskStateStore) Set(
_, err = s.client.comm.Communicate(ctx, genmodels.SetTaskStateStore{
TIID: s.client.tiID,
Key: key,
- Value: value,
+ Value: encoded,
ExpiresAt: expiry,
})
return err
@@ -530,26 +542,49 @@ func resolveExpiry(retention *time.Duration, now
time.Time) (any, error) {
// validateJSONRepresentable checks what the frame encoder actually emits, not
// the Go value: a reflection walk has to mirror the encoder's field rules
(tags,
// "-", omitempty, embedding, marshalers) and misjudges values wherever it
drifts.
-func validateJSONRepresentable(value any) error {
+// On success it returns the encoded bytes, so the caller sends them instead of
+// encoding the value a second time, and whether the value encodes to null.
+func validateJSONRepresentable(value any) (msgpack.RawMessage, bool, error) {
var buf bytes.Buffer
if err := newFrameEncoder(&buf).Encode(value); err != nil {
- return fmt.Errorf("%T is not JSON representable: %w", value,
err)
+ return nil, false, fmt.Errorf("%T is not JSON representable:
%w", value, err)
}
- dec := msgpack.NewDecoder(&buf)
- // The default map decoder fails opaquely on non-string keys.
- dec.SetMapDecoder(func(d *msgpack.Decoder) (any, error) {
- return d.DecodeUntypedMap()
- })
+ encoded := msgpack.RawMessage(buf.Bytes())
+ // Decoded from its own reader so encoded keeps the whole payload.
+ dec := msgpack.NewDecoder(bytes.NewReader(encoded))
+ dec.SetMapDecoder(decodeStringKeyedMap)
decoded, err := dec.DecodeInterface()
if err != nil {
- return fmt.Errorf("%T is not JSON representable: %w", value,
err)
+ return nil, false, fmt.Errorf("%T is not JSON representable:
%w", value, err)
}
- // Checked after decoding: a typed nil such as a nil *string is not ==
nil,
- // yet it still encodes to null, which the Execution API rejects.
- if decoded == nil {
- return errors.New("value must not be nil")
+ if err := validateDecodedValue(decoded); err != nil {
+ return nil, false, err
+ }
+ return encoded, decoded == nil, nil
+}
+
+// decodeStringKeyedMap replaces msgpack's default map decoder, which panics on
+// a uuid.UUID key.
+func decodeStringKeyedMap(d *msgpack.Decoder) (any, error) {
+ n, err := d.DecodeMapLen()
+ if err != nil || n == -1 {
+ return nil, err
+ }
+ m := make(map[string]any, n)
+ for range n {
+ k, err := d.DecodeInterface()
+ if err != nil {
+ return nil, err
+ }
+ key, ok := k.(string)
+ if !ok {
+ return nil, fmt.Errorf("map keys must be strings, got
%T", k)
+ }
+ if m[key], err = d.DecodeInterface(); err != nil {
+ return nil, err
+ }
}
- return validateDecodedValue(decoded)
+ return m, nil
}
func validateDecodedValue(v any) error {
@@ -578,11 +613,8 @@ func validateDecodedValue(v any) error {
}
}
return nil
- case map[any]any:
- for k, elem := range v {
- if _, ok := k.(string); !ok {
- return fmt.Errorf("map keys must be strings,
got %T", k)
- }
+ case map[string]any:
+ for _, elem := range v {
if err := validateDecodedValue(elem); err != nil {
return err
}
diff --git a/go-sdk/pkg/execution/client_test.go
b/go-sdk/pkg/execution/client_test.go
index 2db1a53f769..3c293bfae72 100644
--- a/go-sdk/pkg/execution/client_test.go
+++ b/go-sdk/pkg/execution/client_test.go
@@ -29,6 +29,7 @@ import (
"testing"
"time"
+ "github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -397,6 +398,81 @@ func TestCoordinatorClientPushXComMapIndex(t *testing.T) {
}
}
+// The supervisor drops a SetXCom it cannot decode without replying, so a value
+// that is not JSON must be rejected before it is sent or the task hangs.
+func TestCoordinatorClientPushXComRejectsNonJSONValues(t *testing.T) {
+ tests := []struct {
+ name string
+ value any
+ wantErr string
+ }{
+ {
+ name: "datetime",
+ value: time.Date(2026, 5, 15, 0, 0, 0, 0, time.UTC),
+ wantErr: "time.Time is not JSON representable",
+ },
+ {name: "uuid", value: uuid.New(), wantErr: "[]byte is not JSON
representable"},
+ {name: "byte slice", value: []byte("raw"), wantErr: "[]byte is
not JSON representable"},
+ {
+ // Regression: a key that decodes to a slice used to
panic the decoder.
+ name: "uuid map key",
+ value: map[uuid.UUID]string{uuid.New(): "a"},
+ wantErr: "map keys must be strings",
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ var requestBuf bytes.Buffer
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ client := NewCoordinatorClient(
+ NewCoordinatorComm(&bytes.Buffer{},
&requestBuf, logger), testTIID,
+ )
+ ti := sdk.TaskInstance{DagID: "d", RunID: "r", TaskID:
"t"}
+
+ err := client.PushXCom(context.Background(), ti,
"return_value", tc.value)
+
+ require.ErrorContains(t, err, tc.wantErr)
+ assert.Zero(t, requestBuf.Len(), "a rejected value must
not reach the supervisor")
+ })
+ }
+}
+
+// A task that returns a nil slice or a nil any pushes null, so PushXCom must
not
+// inherit the task state store's nil rule.
+func TestCoordinatorClientPushXComAllowsNil(t *testing.T) {
+ tests := []struct {
+ name string
+ value any
+ }{
+ {name: "nil", value: nil},
+ {name: "typed nil", value: []string(nil)},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ responsePayload := encodeResponseFrame(t, 0,
map[string]any{"type": "OKResponse"}, nil)
+ var responseBuf bytes.Buffer
+ require.NoError(t, writeFrame(&responseBuf,
responsePayload))
+
+ var requestBuf bytes.Buffer
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ client := NewCoordinatorClient(
+ NewCoordinatorComm(&responseBuf, &requestBuf,
logger), testTIID,
+ )
+ ti := sdk.TaskInstance{DagID: "d", RunID: "r", TaskID:
"t"}
+
+ require.NoError(t,
client.PushXCom(context.Background(), ti, "return_value", tc.value))
+
+ sent, err := readFrame(&requestBuf)
+ require.NoError(t, err)
+ sentMap := rawToMap(t, sent.Body)
+ require.Contains(t, sentMap, "value")
+ assert.Nil(t, sentMap["value"])
+ })
+ }
+}
+
// TestCoordinatorClientGetXComMapIndex verifies the GetXCom frame omits
// map_index when none is supplied and propagates a real index otherwise,
// including index 0: a plain deref into the nullable interface{} field would
let
diff --git a/go-sdk/sdk/sdk.go b/go-sdk/sdk/sdk.go
index 75d170c3eb8..1c3f47c8f38 100644
--- a/go-sdk/sdk/sdk.go
+++ b/go-sdk/sdk/sdk.go
@@ -126,6 +126,10 @@ type XComClient interface {
) (any, error)
// PushXCom stores value under key for the given task instance ti.
+ //
+ // value must be JSON serializable; nil is allowed and is pushed as
null.
+ // A time.Time, a byte slice or array such as a uuid.UUID, or a
non-finite
+ // float is rejected before it is sent.
PushXCom(ctx context.Context, ti TaskInstance, key string, value any)
error
}
