From 459edf677e7aa016a982a829a276c0a44bfe4137 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 22 May 2026 17:10:05 +0800 Subject: [PATCH 1/5] go-sdk: Implement coordinator-mode runtime entry point and task runner Wire the supervisor-launched runtime that speaks ADR 0003's coordinator protocol. execution.Serve dials the comm and logs sockets the supervisor passes via the new --comm/--logs flags, installs SocketLogHandler so slog records reach the supervisor, reads StartupDetails, and drives a single TaskInstance through task_runner.Run. The runner injects a CoordinatorClient into the user task function via sdkcontext.SdkClientContextKey so tasks written against the existing sdk.Client API run unchanged. bundlev1server.Serve grows a mode selector so the same binary still serves go-plugin when no coordinator flags are present, and exits non-zero on partial --comm/--logs misuse. DAG-file parsing is intentionally not part of this stack -- it will land in a follow-up once the parsing protocol settles. --- .../bundle/bundlev1/bundlev1server/server.go | 128 +++++++-- go-sdk/example/bundle/main.go | 14 +- go-sdk/pkg/execution/integration_test.go | 250 ++++++++++++++++++ go-sdk/pkg/execution/server.go | 159 +++++++++++ go-sdk/pkg/execution/task_runner.go | 125 +++++++++ 5 files changed, 652 insertions(+), 24 deletions(-) create mode 100644 go-sdk/pkg/execution/integration_test.go create mode 100644 go-sdk/pkg/execution/server.go create mode 100644 go-sdk/pkg/execution/task_runner.go diff --git a/go-sdk/bundle/bundlev1/bundlev1server/server.go b/go-sdk/bundle/bundlev1/bundlev1server/server.go index 67d212ff06976..02036742d0919 100644 --- a/go-sdk/bundle/bundlev1/bundlev1server/server.go +++ b/go-sdk/bundle/bundlev1/bundlev1server/server.go @@ -19,6 +19,7 @@ package bundlev1server import ( "encoding/json" + "errors" "fmt" "log/slog" "os" @@ -32,9 +33,34 @@ import ( "github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server/impl" "github.com/apache/airflow/go-sdk/pkg/bundles/shared" "github.com/apache/airflow/go-sdk/pkg/config" + "github.com/apache/airflow/go-sdk/pkg/execution" ) -var versionInfo *bool = flag.Bool("bundle-metadata", false, "show the embedded bundle info") +// ErrCoordinatorFlagsIncomplete is returned by [Serve] when exactly one of +// --comm or --logs is supplied. Both flags select coordinator mode and must +// be set together; callers (typically main) can check for this sentinel to +// print usage before exiting non-zero. +var ErrCoordinatorFlagsIncomplete = errors.New( + "--comm and --logs must be supplied together", +) + +// CLI Flags. +// The --bundle-metadata flag is used for showing the embedded bundle info in airflow-metadata.yaml spec format. +// The --comm and --logs select the coordinator-mode protocol +// All three are read by Serve to choose a server mode below. +var ( + versionInfo = flag.Bool("bundle-metadata", false, "show the embedded bundle info") + commAddr = flag.String( + "comm", + "", + "host:port of the supervisor's coordinator comm channel (selects coordinator mode)", + ) + logsAddr = flag.String( + "logs", + "", + "host:port of the supervisor's coordinator logs channel (selects coordinator mode)", + ) +) // ServeOpt is an interface for defining options that can be passed to the // Serve function. Each implementation modifies the ServeConfig being @@ -52,24 +78,35 @@ func (s serveConfigFunc) ApplyServeOpt(in *ServerConfig) error { type ServerConfig struct{} -// Serve is the entrypoint for your bundle, and sets it up ready for Airflow's Go Worker to use +// serveMode tags the protocol the binary will speak this run. +type serveMode int + +const ( + modePlugin serveMode = iota // go-plugin gRPC (existing Edge Worker path) + modeMetadataDump // --bundle-metadata: print BundleInfo JSON + modeCoordinator // --comm/--logs: msgpack-over-IPC (ADR 0003) + modeCoordinatorUsageError // misuse: print usage and exit non-zero +) + +// Serve is the entrypoint for your bundle, and sets it up ready for Airflow's +// Go Worker (go-plugin) or Python supervisor (coordinator protocol) to use. // -// Zero or more options to configure the server may also be passed. There are no options yet, this is to allow -// future changes without breaking compatibility +// The mode is decided from CLI flags and process environment. Callers should +// surface the returned error so misuse (e.g. only one of --comm/--logs +// supplied) produces a non-zero exit: +// +// func main() { +// if err := bundlev1server.Serve(&myBundle{}); err != nil { +// log.Fatal(err) +// } +// } +// +// Zero or more options to configure the server may also be passed. There are +// no options yet; the parameter exists to allow future additions without +// breaking compatibility. func Serve(bundle bundlev1.BundleProvider, opts ...ServeOpt) error { config.SetupViper("") - hcLogger := hclog.New(&hclog.LoggerOptions{ - Level: hclog.Trace, - Output: os.Stderr, - JSONFormat: true, - IncludeLocation: true, - AdditionalLocationOffset: 3, - }) - - log := slog.New(hclogslog.Adapt(hcLogger)) - slog.SetDefault(log) - flag.Parse() serveConfig := &ServerConfig{} @@ -77,16 +114,63 @@ func Serve(bundle bundlev1.BundleProvider, opts ...ServeOpt) error { c.ApplyServeOpt(serveConfig) } + switch decideMode() { + case modeMetadataDump: + return dumpBundleMetadata(bundle) + case modeCoordinator: + // In coordinator mode the supervisor reads the logs channel for + // structured records, so configuring the hclog/stderr default + // logger here is unnecessary — execution.Serve installs its own + // slog handler against the logs socket before any user code runs. + return execution.Serve(bundle, *commAddr, *logsAddr) + case modePlugin: + installPluginLogger() + return servePlugin(bundle) + case modeCoordinatorUsageError: + return ErrCoordinatorFlagsIncomplete + } + return nil +} + +func decideMode() serveMode { if *versionInfo { - meta := bundle.GetBundleVersion() - data, err := json.MarshalIndent(meta, "", " ") - if err != nil { - return err - } - fmt.Println(string(data)) - return nil + return modeMetadataDump + } + commSet := *commAddr != "" + logsSet := *logsAddr != "" + if commSet && logsSet { + return modeCoordinator + } + if commSet || logsSet { + // Partial use is a hard error, both flags are required + return modeCoordinatorUsageError + } + return modePlugin +} + +func dumpBundleMetadata(bundle bundlev1.BundleProvider) error { + meta := bundle.GetBundleVersion() + data, err := json.MarshalIndent(meta, "", " ") + if err != nil { + return err } + fmt.Println(string(data)) + return nil +} + +func installPluginLogger() { + hcLogger := hclog.New(&hclog.LoggerOptions{ + Level: hclog.Trace, + Output: os.Stderr, + JSONFormat: true, + IncludeLocation: true, + AdditionalLocationOffset: 3, + }) + log := slog.New(hclogslog.Adapt(hcLogger)) + slog.SetDefault(log) +} +func servePlugin(bundle bundlev1.BundleProvider) error { pluginConfig := &plugin.ServeConfig{ HandshakeConfig: shared.Handshake, Plugins: plugin.PluginSet{ diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 5e970da54951b..350c101b8b0b6 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -20,6 +20,7 @@ package main import ( "context" "fmt" + "log" "log/slog" "runtime" "time" @@ -54,7 +55,9 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { } func main() { - bundlev1server.Serve(&myBundle{}) + if err := bundlev1server.Serve(&myBundle{}); err != nil { + log.Fatal(err) + } } func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, error) { @@ -63,7 +66,14 @@ func extract(ctx context.Context, client sdk.Client, log *slog.Logger) (any, err if err != nil { log.ErrorContext(ctx, "unable to get conn", "error", err) } else { - log.InfoContext(ctx, "got conn", "conn", conn) + // Log only non-sensitive fields; conn.Password and any secrets in + // conn.Extra must never reach the log stream. + log.InfoContext(ctx, "got conn", + "conn_id", conn.ID, + "conn_type", conn.Type, + "host", conn.Host, + "port", conn.Port, + ) } for range 10 { diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go new file mode 100644 index 0000000000000..f251567be21dc --- /dev/null +++ b/go-sdk/pkg/execution/integration_test.go @@ -0,0 +1,250 @@ +// 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 execution + +import ( + "bytes" + "errors" + "io" + "log/slog" + "net" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// --- Test task functions --- + +func failingTask() error { + return errors.New("task failed intentionally") +} + +func panicTask() error { + panic("something went wrong") +} + +func simpleTask() error { + return nil +} + +// buildBundle wires a bundlev1.Registry from a closure and returns it as a +// bundlev1.Bundle (the materialised registry). +func buildBundle(t *testing.T, register func(bundlev1.Registry)) bundlev1.Bundle { + t.Helper() + reg := bundlev1.New() + register(reg) + return reg +} + +// --- Tests --- + +func TestTaskRunnerSuccess(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTask(simpleTask) + }) + + details := &StartupDetails{ + TI: TaskInstanceInfo{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "simpleTask", + RunID: "run1", + MapIndex: -1, + }, + BundleInfo: BundleInfoMsg{Name: "test", Version: "1.0"}, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(bundle, details, comm, logger) + assert.Equal(t, "SucceedTask", result["type"]) +} + +func TestTaskRunnerFailure(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTask(failingTask) + }) + + details := &StartupDetails{ + TI: TaskInstanceInfo{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "failingTask", + RunID: "run1", + MapIndex: -1, + }, + BundleInfo: BundleInfoMsg{Name: "test", Version: "1.0"}, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(bundle, details, comm, logger) + assert.Equal(t, "TaskState", result["type"]) + assert.Equal(t, "failed", result["state"]) +} + +func TestTaskRunnerTaskNotFound(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTask(simpleTask) + }) + + details := &StartupDetails{ + TI: TaskInstanceInfo{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "nonexistent", + RunID: "run1", + }, + BundleInfo: BundleInfoMsg{Name: "test", Version: "1.0"}, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(bundle, details, comm, logger) + assert.Equal(t, "TaskState", result["type"]) + assert.Equal(t, "removed", result["state"]) +} + +func TestTaskRunnerPanic(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTask(panicTask) + }) + + details := &StartupDetails{ + TI: TaskInstanceInfo{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "panicTask", + RunID: "run1", + MapIndex: -1, + }, + BundleInfo: BundleInfoMsg{Name: "test", Version: "1.0"}, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(bundle, details, comm, logger) + assert.Equal(t, "TaskState", result["type"]) + assert.Equal(t, "failed", result["state"]) +} + +// --- End-to-end Serve test against a fake supervisor --- + +// fakeProvider implements bundlev1.BundleProvider; it lets a test inject the +// registration closure and a synthetic version. +type fakeProvider struct { + register func(bundlev1.Registry) error +} + +func (f *fakeProvider) GetBundleVersion() bundlev1.BundleInfo { + v := "1.0" + return bundlev1.BundleInfo{Name: "fake", Version: &v} +} + +func (f *fakeProvider) RegisterDags(reg bundlev1.Registry) error { + if f.register == nil { + return nil + } + return f.register(reg) +} + +func startSupervisor( + t *testing.T, +) (commAddr, logsAddr string, commCh, logsCh chan net.Conn, cleanup func()) { + t.Helper() + commLn, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + logsLn, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + commCh = make(chan net.Conn, 1) + logsCh = make(chan net.Conn, 1) + go func() { + c, err := commLn.Accept() + if err == nil { + commCh <- c + } + close(commCh) + }() + go func() { + c, err := logsLn.Accept() + if err == nil { + logsCh <- c + } + close(logsCh) + }() + cleanup = func() { + commLn.Close() + logsLn.Close() + } + return commLn.Addr().String(), logsLn.Addr().String(), commCh, logsCh, cleanup +} + +func TestServeStartupDetailsEndToEnd(t *testing.T) { + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + defer cleanup() + + provider := &fakeProvider{ + register: func(r bundlev1.Registry) error { + r.AddDag("dag1").AddTask(simpleTask) + return nil + }, + } + + done := make(chan error, 1) + go func() { done <- Serve(provider, commAddr, logsAddr) }() + + commConn := <-commCh + defer commConn.Close() + logsConn := <-logsCh + defer logsConn.Close() + + payload, err := encodeRequest(0, map[string]any{ + "type": "StartupDetails", + "ti": map[string]any{ + "id": "550e8400-e29b-41d4-a716-446655440000", + "dag_id": "dag1", + "task_id": "simpleTask", + "run_id": "run1", + "try_number": 1, + }, + "bundle_info": map[string]any{"name": "fake", "version": "1.0"}, + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, payload)) + + frame, err := readFrame(commConn) + require.NoError(t, err) + require.Nil(t, frame.Err) + assert.Equal(t, "SucceedTask", frame.Body["type"]) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after task completion") + } +} diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go new file mode 100644 index 0000000000000..96340f79f21c6 --- /dev/null +++ b/go-sdk/pkg/execution/server.go @@ -0,0 +1,159 @@ +// 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 execution implements the SDK coordinator-protocol runtime +// (msgpack-over-IPC). It is the second mode of bundlev1server.Serve: when +// the bundle binary is launched with --comm/--logs by the Airflow supervisor +// (Python ExecutableCoordinator), bundlev1server.Serve dispatches here. +// +// The first inbound frame on the comm socket is a StartupDetails message +// that drives multi-round task execution. +// +// See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md. +package execution + +import ( + "fmt" + "log/slog" + "net" + "sync" + "time" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" +) + +// dialTimeout bounds how long execution.Serve waits to reach the supervisor's +// comm and logs sockets. The supervisor opens the listeners before spawning +// the bundle binary, so the dials normally succeed in milliseconds; the +// timeout exists so an unreachable address fails fast instead of hanging the +// runtime indefinitely. +const dialTimeout = 30 * time.Second + +// Serve runs the bundle binary in coordinator mode. It dials the supervisor's +// comm and logs sockets, installs an slog handler that writes JSON-line +// records to the logs connection, and dispatches on the first frame. +// +// Serve returns nil on a clean shutdown (task execution completed); a +// non-nil error indicates a protocol-level failure (connection loss, +// malformed frames, unknown first message type). +func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { + if commAddr == "" { + return fmt.Errorf("missing --comm=host:port argument") + } + if logsAddr == "" { + return fmt.Errorf("missing --logs=host:port argument") + } + + // Buffer log records until the logs socket is connected. Anything the + // runtime emits between Connect-time and the first frame still gets + // flushed. + logHandler := NewSocketLogHandler(nil, slog.LevelDebug) + logger := slog.New(logHandler) + slog.SetDefault(logger) + + // Connect to both sockets concurrently so the supervisor can accept them + // in either order. + dialer := &net.Dialer{Timeout: dialTimeout} + var commConn, logsConn net.Conn + var commErr, logsErr error + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + commConn, commErr = dialer.Dial("tcp", commAddr) + }() + go func() { + defer wg.Done() + logsConn, logsErr = dialer.Dial("tcp", logsAddr) + }() + wg.Wait() + + // Either dial may succeed while the other fails; close any orphaned + // connection before returning so we don't leak an open TCP socket. + if commErr != nil { + if logsConn != nil { + logsConn.Close() + } + return fmt.Errorf("connecting to comm socket %s: %w", commAddr, commErr) + } + if logsErr != nil { + commConn.Close() + return fmt.Errorf("connecting to logs socket %s: %w", logsAddr, logsErr) + } + defer commConn.Close() + defer logsConn.Close() + + logHandler.Connect(logsConn) + logger.Debug("Connected", "comm", commAddr, "logs", logsAddr) + + // Materialise the bundle (RegisterDags) up front. Both protocol paths + // need the registry, and doing it once before the first frame keeps the + // dispatcher simple. + bundle, err := materialiseBundle(provider) + if err != nil { + return fmt.Errorf("registering dags: %w", err) + } + + comm := NewCoordinatorComm(commConn, commConn, logger) + + frame, err := comm.ReadMessage() + if err != nil { + return fmt.Errorf("reading initial message: %w", err) + } + + if frame.Err != nil { + errResp := decodeErrorResponse(frame.Err) + if errResp != nil { + return fmt.Errorf( + "received error from supervisor: [%s] %v", + errResp.Error, + errResp.Detail, + ) + } + } + + body, err := decodeIncomingBody(frame.Body) + if err != nil { + return fmt.Errorf("decoding initial message: %w", err) + } + + switch msg := body.(type) { + case *StartupDetails: + logger.Debug("Task execution mode", + "dag_id", msg.TI.DagID, + "task_id", msg.TI.TaskID, + ) + result := RunTask(bundle, msg, comm, logger) + if err := comm.SendRequest(frame.ID, result); err != nil { + return fmt.Errorf("sending task result: %w", err) + } + logger.Debug("Task execution complete") + + default: + return fmt.Errorf("unexpected initial message type: %T", body) + } + + return nil +} + +func materialiseBundle(provider bundlev1.BundleProvider) (bundlev1.Bundle, error) { + reg := bundlev1.New() + if err := provider.RegisterDags(reg); err != nil { + return nil, err + } + return reg, nil +} diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go new file mode 100644 index 0000000000000..df6f37c5f2712 --- /dev/null +++ b/go-sdk/pkg/execution/task_runner.go @@ -0,0 +1,125 @@ +// 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 execution + +import ( + "context" + "log/slog" + "runtime/debug" + "time" + + "github.com/google/uuid" + + "github.com/apache/airflow/go-sdk/bundle/bundlev1" + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +// RunTask executes a task based on StartupDetails received from the supervisor. +// +// It looks up the task in the bundle, creates a CoordinatorClient for SDK +// calls, executes the task, and returns a terminal message body +// (SucceedTaskMsg or TaskStateMsg) ready to ship as the final response frame. +// +// The supervisor owns the Execution-API state transitions in coordinator +// mode, so we deliberately bypass worker.ExecuteTaskWorkload (which drives +// Run / UpdateState itself) and only invoke the user's task function. +func RunTask( + bundle bundlev1.Bundle, + details *StartupDetails, + comm *CoordinatorComm, + logger *slog.Logger, +) map[string]any { + task, exists := bundle.LookupTask(details.TI.DagID, details.TI.TaskID) + if !exists { + logger.Error("Task not registered", + "dag_id", details.TI.DagID, + "task_id", details.TI.TaskID, + ) + return TaskStateMsg{State: TaskStateRemoved, EndDate: time.Now().UTC()}.toMap() + } + + client := NewCoordinatorClient(comm) + + // taskFunction.sendXcom reads the workload from context to get the task + // instance ids; populate it the same shape the gRPC path uses. + tiUUID, err := uuid.Parse(details.TI.ID) + if err != nil { + logger.Error("Invalid task instance UUID from supervisor", + "dag_id", details.TI.DagID, + "task_id", details.TI.TaskID, + "ti_id", details.TI.ID, + "error", err, + ) + return TaskStateMsg{State: TaskStateFailed, EndDate: time.Now().UTC()}.toMap() + } + mapIndex := details.TI.MapIndex + workload := api.ExecuteTaskWorkload{ + TI: api.TaskInstance{ + Id: tiUUID, + DagId: details.TI.DagID, + RunId: details.TI.RunID, + TaskId: details.TI.TaskID, + TryNumber: details.TI.TryNumber, + MapIndex: &mapIndex, + }, + BundleInfo: api.BundleInfo{ + Name: details.BundleInfo.Name, + Version: &details.BundleInfo.Version, + }, + } + + ctx := context.Background() + ctx = context.WithValue(ctx, sdkcontext.WorkloadContextKey, workload) + ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) + + return executeTask(ctx, task, logger) +} + +// executeTask runs the task and handles success, failure, and panics. +func executeTask( + ctx context.Context, + task bundlev1.Task, + logger *slog.Logger, +) (result map[string]any) { + defer func() { + if r := recover(); r != nil { + logger.Error("Recovered panic in task", + "error", r, + "stack", string(debug.Stack()), + ) + result = TaskStateMsg{ + State: TaskStateFailed, + EndDate: time.Now().UTC(), + }.toMap() + } + }() + + if err := task.Execute(ctx, logger); err != nil { + logger.ErrorContext(ctx, "Task failed", "error", err) + return TaskStateMsg{ + State: TaskStateFailed, + EndDate: time.Now().UTC(), + }.toMap() + } + + return SucceedTaskMsg{ + EndDate: time.Now().UTC(), + }.toMap() +} From 806c14fef10c9764c55f3dddc1d34177747f4f13 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 31 May 2026 16:10:12 +0800 Subject: [PATCH 2/5] Add coordinator-mode client round-trip integration test The existing end-to-end test only drove a no-op task, so the comm dispatcher's request/response multiplexing -- the core of coordinator mode, where a running task calls back into the supervisor mid-execution -- was never exercised against the real Serve entry point. Add a test that runs Serve against a fake supervisor over a loopback socket pair and registers a task that pulls a variable (GetVariable) and returns a value (triggering a return-value SetXCom push). The fake supervisor must answer both runtime-initiated requests, on distinct frame ids, before the terminal SucceedTask frame arrives, covering the full multiplexed round trip. A connection deadline keeps a regression (e.g. the env-var fast path swallowing the request, or a dispatcher deadlock) from hanging until the Go test timeout. --- go-sdk/pkg/execution/integration_test.go | 107 +++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index f251567be21dc..7279899d89266 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -19,6 +19,7 @@ package execution import ( "bytes" + "context" "errors" "io" "log/slog" @@ -30,6 +31,7 @@ import ( "github.com/stretchr/testify/require" "github.com/apache/airflow/go-sdk/bundle/bundlev1" + "github.com/apache/airflow/go-sdk/sdk" ) // --- Test task functions --- @@ -248,3 +250,108 @@ func TestServeStartupDetailsEndToEnd(t *testing.T) { t.Fatal("Serve did not return after task completion") } } + +// TestServeClientRoundTripEndToEnd drives a task that calls back into the +// supervisor mid-execution, so the comm dispatcher's request/response +// multiplexing is exercised against the real Serve rather than only the +// no-op task path. The registered task pulls a variable (GetVariable) and +// returns a value (which triggers a return-value SetXCom push); the fake +// supervisor must answer both runtime-initiated requests before the terminal +// SucceedTask frame is sent. +func TestServeClientRoundTripEndToEnd(t *testing.T) { + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + defer cleanup() + + // Unique key so the GetVariable env-var fast path + // (AIRFLOW_VAR_) cannot short-circuit the socket round trip. + const varKey = "go_sdk_round_trip_only_key" + + var gotVar string + provider := &fakeProvider{ + register: func(r bundlev1.Registry) error { + r.AddDag("dag1").AddTaskWithName("getvar", + func(ctx context.Context, c sdk.Client) (string, error) { + v, err := c.GetVariable(ctx, varKey) + if err != nil { + return "", err + } + gotVar = v + return "xval", nil + }) + return nil + }, + } + + done := make(chan error, 1) + go func() { done <- Serve(provider, commAddr, logsAddr) }() + + commConn := <-commCh + defer commConn.Close() + logsConn := <-logsCh + defer logsConn.Close() + + // Bound every read/write so a regression (e.g. the env-var fast path + // swallowing the request, or a dispatcher deadlock) fails fast instead of + // hanging until the Go test timeout. + require.NoError(t, commConn.SetDeadline(time.Now().Add(10*time.Second))) + + // 1. Kick off task execution. + startup, err := encodeRequest(0, map[string]any{ + "type": "StartupDetails", + "ti": map[string]any{ + "id": "550e8400-e29b-41d4-a716-446655440000", + "dag_id": "dag1", + "task_id": "getvar", + "run_id": "run1", + "try_number": 1, + }, + "bundle_info": map[string]any{"name": "fake", "version": "1.0"}, + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, startup)) + + // 2. The task's GetVariable call blocks until the supervisor answers. + varReq, err := readFrame(commConn) + require.NoError(t, err) + require.Nil(t, varReq.Err) + assert.Equal(t, "GetVariable", varReq.Body["type"]) + assert.Equal(t, varKey, varReq.Body["key"]) + + varReply, err := encodeRequest(varReq.ID, map[string]any{ + "type": "VariableResult", + "key": varKey, + "value": "hello", + }) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, varReply)) + + // 3. Returning a value triggers a return-value XCom push; answer it with + // an empty (non-error) response so PushXCom unblocks. + xcomReq, err := readFrame(commConn) + require.NoError(t, err) + require.Nil(t, xcomReq.Err) + assert.Equal(t, "SetXCom", xcomReq.Body["type"]) + assert.Equal(t, "return_value", xcomReq.Body["key"]) + assert.Equal(t, "xval", xcomReq.Body["value"]) + assert.NotEqual(t, varReq.ID, xcomReq.ID, "second runtime request must use a fresh frame id") + + xcomReply, err := encodeRequest(xcomReq.ID, map[string]any{}) + require.NoError(t, err) + require.NoError(t, writeFrame(commConn, xcomReply)) + + // 4. With both calls answered, the task finishes and Serve ships the + // terminal SucceedTask frame on the StartupDetails frame id. + term, err := readFrame(commConn) + require.NoError(t, err) + require.Nil(t, term.Err) + assert.Equal(t, "SucceedTask", term.Body["type"]) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after task completion") + } + + assert.Equal(t, "hello", gotVar) +} From 992a7bab2adfaeff40f4224b85464fc182b8abc0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 31 May 2026 16:12:12 +0800 Subject: [PATCH 3/5] Log coordinator-mode runtime failures and document exit contract When the coordinator runtime fails after connecting to the supervisor (bundle registration error, an unreadable or undecodable initial frame, a supervisor-reported error, or an unexpected first message type) it returned a bare error and relied on the process exit code alone to signal failure. The supervisor only sees the closed comm socket, with no explanation in its log stream, making such failures hard to diagnose. Log the reason at Error on each post-connect failure path so it reaches the supervisor over the already-connected logs socket, and document the failure-signaling contract on Serve: a non-nil return must become a non-zero process exit, which the supervisor records as FAILED (or UP_FOR_RETRY), and a structured TaskState frame is only honored when the process exits 0. This makes the fail-closed behavior explicit so a later refactor cannot accidentally turn a startup failure into a success. Adds a test that a RegisterDags failure makes Serve return the error without writing a terminal frame, leaving the supervisor to observe the comm socket closing. --- go-sdk/pkg/execution/integration_test.go | 37 ++++++++++++++++++++++++ go-sdk/pkg/execution/server.go | 24 +++++++++++++-- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 7279899d89266..5c9b8e52ac755 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -355,3 +355,40 @@ func TestServeClientRoundTripEndToEnd(t *testing.T) { assert.Equal(t, "hello", gotVar) } + +// TestServeRegisterDagsFailureClosesComm asserts the failure-signaling +// contract: when bundle registration fails after the sockets are connected, +// Serve returns the error (so the caller exits non-zero) without writing a +// terminal frame. The supervisor observes the failure as the comm socket +// closing rather than as a TaskState message. +func TestServeRegisterDagsFailureClosesComm(t *testing.T) { + commAddr, logsAddr, commCh, logsCh, cleanup := startSupervisor(t) + defer cleanup() + + wantErr := errors.New("boom registering dags") + provider := &fakeProvider{ + register: func(bundlev1.Registry) error { return wantErr }, + } + + done := make(chan error, 1) + go func() { done <- Serve(provider, commAddr, logsAddr) }() + + commConn := <-commCh + defer commConn.Close() + logsConn := <-logsCh + defer logsConn.Close() + + select { + case err := <-done: + require.Error(t, err) + assert.ErrorIs(t, err, wantErr) + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return after RegisterDags failure") + } + + // No terminal frame was sent: the next read on the comm socket sees the + // connection close instead of a decodable frame. + require.NoError(t, commConn.SetReadDeadline(time.Now().Add(time.Second))) + _, err := readFrame(commConn) + require.Error(t, err) +} diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index 96340f79f21c6..f99e456d4081b 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -47,9 +47,21 @@ const dialTimeout = 30 * time.Second // comm and logs sockets, installs an slog handler that writes JSON-line // records to the logs connection, and dispatches on the first frame. // -// Serve returns nil on a clean shutdown (task execution completed); a +// Serve returns nil on a clean shutdown: the task ran and its terminal +// TaskState/SucceedTask frame was delivered, and the caller should exit 0. A // non-nil error indicates a protocol-level failure (connection loss, -// malformed frames, unknown first message type). +// malformed frames, unknown first message type) that happens before or +// instead of delivering a terminal frame. +// +// Failure-signaling contract: the caller (main) must turn a non-nil error +// into a non-zero process exit. The supervisor derives the task's final state +// primarily from the child's exit code -- a non-zero exit is recorded as +// FAILED (or UP_FOR_RETRY when retries are configured), and a structured +// TaskState frame is only honored when the process exits 0 (see the Python +// supervisor's ActivitySubprocess.final_state). So an early error return here +// fails closed without needing to send a frame; the post-connect paths below +// log the reason at Error first so it still reaches the supervisor's log +// stream over the already-connected logs socket. func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { if commAddr == "" { return fmt.Errorf("missing --comm=host:port argument") @@ -105,6 +117,7 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { // dispatcher simple. bundle, err := materialiseBundle(provider) if err != nil { + logger.Error("Bundle registration failed", "error", err) return fmt.Errorf("registering dags: %w", err) } @@ -112,12 +125,17 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { frame, err := comm.ReadMessage() if err != nil { + logger.Error("Failed to read initial message from supervisor", "error", err) return fmt.Errorf("reading initial message: %w", err) } if frame.Err != nil { errResp := decodeErrorResponse(frame.Err) if errResp != nil { + logger.Error("Supervisor reported an error on the initial frame", + "error", errResp.Error, + "detail", errResp.Detail, + ) return fmt.Errorf( "received error from supervisor: [%s] %v", errResp.Error, @@ -128,6 +146,7 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { body, err := decodeIncomingBody(frame.Body) if err != nil { + logger.Error("Failed to decode initial message", "error", err) return fmt.Errorf("decoding initial message: %w", err) } @@ -144,6 +163,7 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { logger.Debug("Task execution complete") default: + logger.Error("Unexpected initial message type", "type", fmt.Sprintf("%T", body)) return fmt.Errorf("unexpected initial message type: %T", body) } From c5af270940d8a0234531934242faff308d5ae15f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 31 May 2026 16:14:35 +0800 Subject: [PATCH 4/5] Cancel coordinator-mode tasks on SIGINT/SIGTERM The coordinator runtime ran each task under a fresh context.Background(), so a supervisor shutdown could not reach the task: SIGTERM terminated the process outright and a cooperative task had no chance to observe the shutdown and return. The final-frame write was also unbounded, so a half-open comm socket could wedge the runtime indefinitely. Derive the task's root context from signal.NotifyContext(SIGINT, SIGTERM) in Serve and thread it through RunTask into the user task, so a ctx-aware task returns promptly on a supervisor shutdown; a task that ignores ctx is still stopped by the supervisor's follow-up SIGKILL. Bound the terminal frame write with a deadline so a wedged socket fails fast (non-zero exit) instead of hanging. An in-flight client call is unaffected: it already unblocks on supervisor disconnect via the comm dispatcher's ErrDispatcherClosed path. Adds a test that a cancelled root context reaches the user task through RunTask and is reported as a failed terminal state. --- go-sdk/pkg/execution/integration_test.go | 38 +++++++++++++++++++++--- go-sdk/pkg/execution/server.go | 22 +++++++++++++- go-sdk/pkg/execution/task_runner.go | 5 +++- 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 5c9b8e52ac755..ed0f437cff63a 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -78,7 +78,7 @@ func TestTaskRunnerSuccess(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, nil)) comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - result := RunTask(bundle, details, comm, logger) + result := RunTask(context.Background(), bundle, details, comm, logger) assert.Equal(t, "SucceedTask", result["type"]) } @@ -101,7 +101,7 @@ func TestTaskRunnerFailure(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, nil)) comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - result := RunTask(bundle, details, comm, logger) + result := RunTask(context.Background(), bundle, details, comm, logger) assert.Equal(t, "TaskState", result["type"]) assert.Equal(t, "failed", result["state"]) } @@ -124,7 +124,7 @@ func TestTaskRunnerTaskNotFound(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, nil)) comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - result := RunTask(bundle, details, comm, logger) + result := RunTask(context.Background(), bundle, details, comm, logger) assert.Equal(t, "TaskState", result["type"]) assert.Equal(t, "removed", result["state"]) } @@ -148,7 +148,37 @@ func TestTaskRunnerPanic(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, nil)) comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - result := RunTask(bundle, details, comm, logger) + result := RunTask(context.Background(), bundle, details, comm, logger) + assert.Equal(t, "TaskState", result["type"]) + assert.Equal(t, "failed", result["state"]) +} + +func TestRunTaskHonorsContextCancellation(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("ctxcheck", + func(ctx context.Context) error { return ctx.Err() }) + }) + + details := &StartupDetails{ + TI: TaskInstanceInfo{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "ctxcheck", + RunID: "run1", + MapIndex: -1, + }, + BundleInfo: BundleInfoMsg{Name: "test", Version: "1.0"}, + } + + // A cancelled root context must reach the user task through RunTask's + // threading; the task surfaces ctx.Err(), which RunTask maps to failed. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(ctx, bundle, details, comm, logger) assert.Equal(t, "TaskState", result["type"]) assert.Equal(t, "failed", result["state"]) } diff --git a/go-sdk/pkg/execution/server.go b/go-sdk/pkg/execution/server.go index f99e456d4081b..ade60d2295604 100644 --- a/go-sdk/pkg/execution/server.go +++ b/go-sdk/pkg/execution/server.go @@ -27,10 +27,13 @@ package execution import ( + "context" "fmt" "log/slog" "net" + "os/signal" "sync" + "syscall" "time" "github.com/apache/airflow/go-sdk/bundle/bundlev1" @@ -43,6 +46,13 @@ import ( // runtime indefinitely. const dialTimeout = 30 * time.Second +// terminalSendTimeout bounds the write of the final TaskState/SucceedTask +// frame. The supervisor normally drains the comm socket promptly, but a +// half-open connection (the supervisor gone without a clean close) could +// otherwise wedge the runtime on a blocked write; the deadline turns that +// into a fast failure -- and thus a non-zero exit -- instead of a hang. +const terminalSendTimeout = 30 * time.Second + // Serve runs the bundle binary in coordinator mode. It dials the supervisor's // comm and logs sockets, installs an slog handler that writes JSON-line // records to the logs connection, and dispatches on the first frame. @@ -70,6 +80,14 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { return fmt.Errorf("missing --logs=host:port argument") } + // A supervisor shutdown arrives as SIGTERM (escalated to SIGKILL after a + // grace period). Trap SIGINT/SIGTERM into a context so a cooperative, + // ctx-aware task can observe the shutdown and return promptly. A task that + // ignores ctx is still stopped by the supervisor's follow-up SIGKILL, so + // trapping the signal here does not strand a non-cooperative task. + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + // Buffer log records until the logs socket is connected. Anything the // runtime emits between Connect-time and the first frame still gets // flushed. @@ -156,7 +174,9 @@ func Serve(provider bundlev1.BundleProvider, commAddr, logsAddr string) error { "dag_id", msg.TI.DagID, "task_id", msg.TI.TaskID, ) - result := RunTask(bundle, msg, comm, logger) + result := RunTask(ctx, bundle, msg, comm, logger) + // Bound the terminal write so a wedged socket cannot hang shutdown. + _ = commConn.SetWriteDeadline(time.Now().Add(terminalSendTimeout)) if err := comm.SendRequest(frame.ID, result); err != nil { return fmt.Errorf("sending task result: %w", err) } diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index df6f37c5f2712..4eda21a523e4f 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -40,7 +40,11 @@ import ( // The supervisor owns the Execution-API state transitions in coordinator // mode, so we deliberately bypass worker.ExecuteTaskWorkload (which drives // Run / UpdateState itself) and only invoke the user's task function. +// +// ctx is the task's root context; Serve derives it from SIGINT/SIGTERM, so a +// cooperative task that honors ctx returns promptly on a supervisor shutdown. func RunTask( + ctx context.Context, bundle bundlev1.Bundle, details *StartupDetails, comm *CoordinatorComm, @@ -85,7 +89,6 @@ func RunTask( }, } - ctx := context.Background() ctx = context.WithValue(ctx, sdkcontext.WorkloadContextKey, workload) ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) From 66734ab0263ef1a6b7c8adac2a6a8542be0ca8b0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 31 May 2026 16:53:41 +0800 Subject: [PATCH 5/5] Track up for retry transition as TODO --- go-sdk/pkg/execution/task_runner.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index 4eda21a523e4f..dba184b588e54 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -116,6 +116,10 @@ func executeTask( if err := task.Execute(ctx, logger); err != nil { logger.ErrorContext(ctx, "Task failed", "error", err) + // TODO(https://github.com/apache/airflow/issues/67797): emit RetryTask + // (UP_FOR_RETRY) when ti_context.should_retry is set. Today every + // failure maps to terminal FAILED because the supervisor honors this + // frame on exit 0 and we never send RetryTask, so retries are lost. return TaskStateMsg{ State: TaskStateFailed, EndDate: time.Now().UTC(),