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..ed0f437cff63a --- /dev/null +++ b/go-sdk/pkg/execution/integration_test.go @@ -0,0 +1,424 @@ +// 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" + "context" + "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" + "github.com/apache/airflow/go-sdk/sdk" +) + +// --- 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(context.Background(), 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(context.Background(), 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(context.Background(), 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(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"]) +} + +// --- 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") + } +} + +// 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) +} + +// 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 new file mode 100644 index 0000000000000..ade60d2295604 --- /dev/null +++ b/go-sdk/pkg/execution/server.go @@ -0,0 +1,199 @@ +// 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 ( + "context" + "fmt" + "log/slog" + "net" + "os/signal" + "sync" + "syscall" + "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 + +// 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. +// +// 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) 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") + } + if logsAddr == "" { + 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. + 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 { + logger.Error("Bundle registration failed", "error", err) + return fmt.Errorf("registering dags: %w", err) + } + + comm := NewCoordinatorComm(commConn, commConn, logger) + + 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, + errResp.Detail, + ) + } + } + + 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) + } + + switch msg := body.(type) { + case *StartupDetails: + logger.Debug("Task execution mode", + "dag_id", msg.TI.DagID, + "task_id", msg.TI.TaskID, + ) + 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) + } + 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) + } + + 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..dba184b588e54 --- /dev/null +++ b/go-sdk/pkg/execution/task_runner.go @@ -0,0 +1,132 @@ +// 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. +// +// 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, + 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.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) + // 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(), + }.toMap() + } + + return SucceedTaskMsg{ + EndDate: time.Now().UTC(), + }.toMap() +}