You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jason810496 (via GitHub)" <gi...@apache.org> on 2026/05/19 06:10:24 UTC
[PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
jason810496 opened a new pull request, #67154:
URL: https://github.com/apache/airflow/pull/67154
related: ADR 0003 (coordinator-protocol msgpack-over-IPC) — merged in a separate PR
## Why
The Go SDK bundle binary currently only speaks the go-plugin protocol
used by the Edge Worker. To let the same binary run under the
ExecutableCoordinator path on the Python side (per ADR 0003), the
bundle needs to be able to act as a coordinator-protocol server:
accept startup details from the supervisor over stdio, drive a task
through its lifecycle, and stream logs / comms frames back. This PR
adds that runtime as a new `pkg/execution` package and wires the
bundle server to invoke it.
## How
- New `go-sdk/pkg/execution` package implements the coordinator-side
runtime: msgpack framing (`frames.go`), the message catalogue
(`messages.go`), Airflow-shaped serialisation (`serde.go`), the
comms decoder/encoder pair (`comms.go`), an stdio-backed logger
(`logger.go`), a `Client` for outbound supervisor calls
(`client.go`), the `Server` that owns the stdio loop (`server.go`),
and a `TaskRunner` that drives a task from startup-details through
to completion (`task_runner.go`).
- `dag_parser.go` reuses the existing registry to enumerate DAGs/tasks
for the supervisor without spawning the legacy go-plugin path.
- Bundle plumbing: `bundlev1server/server.go`, `bundlev1/registry.go`,
and `bundlev1/task.go` learn to hand a task off to the new runner;
`sdkcontext/keys.go` exposes the new context keys; `sdk/client.go`
and `sdk/connection.go` get the small adjustments required for the
new entry path.
## What
- Add `go-sdk/pkg/execution/{client,comms,dag_parser,frames,logger,messages,serde,server,task_runner}.go`
plus matching unit and integration tests.
- Extend `go-sdk/bundle/bundlev1/{bundlev1server/server.go,registry.go,task.go}`
to route coordinator-protocol invocations through the new runtime.
- Add the `sdkcontext` keys consumed by the runner and small SDK
client/connection tweaks.
- Bump `go.mod` / `go.sum` for the new transport/codec dependencies.
---
##### Was generative AI tooling used to co-author this PR?
- [x] Yes, with help of Claude Code Opus 4.7 following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [For review only, Don't merge] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #67154:
URL: https://github.com/apache/airflow/pull/67154#issuecomment-4609622122
All sub PRs just merged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #67154:
URL: https://github.com/apache/airflow/pull/67154#discussion_r3280494326
##########
go-sdk/example/bundle/main.go:
##########
@@ -54,7 +55,9 @@
}
func main() {
- bundlev1server.Serve(&myBundle{})
+ if err := bundlev1server.Serve(&myBundle{}); err != nil {
+ log.Fatal(err)
Review Comment:
## CodeQL / Clear-text logging of sensitive information
[Sensitive data returned by an access to Password](1) flows to a logging call.
[Show more details](https://github.com/apache/airflow/security/code-scanning/604)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #67154:
URL: https://github.com/apache/airflow/pull/67154#discussion_r3279430006
##########
go-sdk/pkg/execution/comms.go:
##########
@@ -0,0 +1,128 @@
+// 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 (
+ "fmt"
+ "io"
+ "log/slog"
+ "sync"
+ "sync/atomic"
+)
+
+// CoordinatorComm manages bidirectional communication with the Airflow supervisor
+// over a length-prefixed msgpack socket connection.
+type CoordinatorComm struct {
+ reader io.Reader
+ writer io.Writer
+ nextID atomic.Int32
+ logger *slog.Logger
+
+ wmu sync.Mutex // serialises writes
+ rmu sync.Mutex // serialises reads
+}
+
+// NewCoordinatorComm creates a new communication channel.
+func NewCoordinatorComm(reader io.Reader, writer io.Writer, logger *slog.Logger) *CoordinatorComm {
+ return &CoordinatorComm{
+ reader: reader,
+ writer: writer,
+ logger: logger,
+ }
+}
+
+// ReadMessage reads and decodes one frame from the comm socket.
+// It returns the raw IncomingFrame with decoded map bodies.
+func (c *CoordinatorComm) ReadMessage() (IncomingFrame, error) {
+ c.rmu.Lock()
+ defer c.rmu.Unlock()
+ frame, err := readFrame(c.reader)
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading frame: %w", err)
+ }
+ c.logger.Debug("Received frame", "id", frame.ID)
+ return frame, nil
+}
+
+// SendRequest sends a request frame (2-element: [id, body]) to the supervisor.
+func (c *CoordinatorComm) SendRequest(id int, body map[string]any) error {
+ payload, err := encodeRequest(id, body)
+ if err != nil {
+ return fmt.Errorf("encoding request: %w", err)
+ }
+ c.logger.Debug("Sending request", "id", id)
+ c.wmu.Lock()
+ defer c.wmu.Unlock()
+ return writeFrame(c.writer, payload)
+}
+
+// Communicate sends a request and waits for the corresponding response.
+// This is a synchronous request-response: the caller sends a request and blocks
+// until the next frame arrives. The protocol is single-threaded on the comm socket.
+//
+// If the response contains an error element, it is returned as an ApiError.
+// Otherwise, the response body map is returned.
+func (c *CoordinatorComm) Communicate(body map[string]any) (map[string]any, error) {
+ id := int(c.nextID.Add(1) - 1)
+
+ if err := c.SendRequest(id, body); err != nil {
+ return nil, err
Review Comment:
Addressed in https://github.com/apache/airflow/pull/67154/commits/a4e91c5aa58c437052806ac7b93cc5d1b053f267
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #67154:
URL: https://github.com/apache/airflow/pull/67154#discussion_r3278889575
##########
go-sdk/pkg/execution/frames.go:
##########
@@ -0,0 +1,229 @@
+// 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"
+ "encoding/binary"
+ "fmt"
+ "io"
+
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// IncomingFrame represents a decoded frame received from the comm socket.
+type IncomingFrame struct {
+ ID int
+ Body map[string]any
+ Err map[string]any // non-nil only for response frames (3-element arrays)
+}
+
+// encodeRequest encodes a request frame (2-element msgpack array: [id, body]).
+func encodeRequest(id int, body map[string]any) ([]byte, error) {
+ var buf bytes.Buffer
+ enc := msgpack.NewEncoder(&buf)
+ enc.UseCompactInts(true)
+
+ if err := enc.EncodeArrayLen(2); err != nil {
+ return nil, err
+ }
+ if err := enc.EncodeInt(int64(id)); err != nil {
+ return nil, err
+ }
+ if err := enc.Encode(body); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+}
+
+// writeFrame writes a length-prefixed msgpack payload to the writer.
+// Format: [4-byte big-endian length][payload bytes]
+func writeFrame(w io.Writer, payload []byte) error {
+ prefix := make([]byte, 4)
+ binary.BigEndian.PutUint32(prefix, uint32(len(payload)))
+ if _, err := w.Write(prefix); err != nil {
+ return fmt.Errorf("writing length prefix: %w", err)
+ }
+ if _, err := w.Write(payload); err != nil {
+ return fmt.Errorf("writing payload: %w", err)
+ }
+ return nil
+}
+
+// readFrame reads one length-prefixed msgpack frame from the reader and decodes it.
+func readFrame(r io.Reader) (IncomingFrame, error) {
+ // Read 4-byte big-endian length prefix.
+ prefix := make([]byte, 4)
+ if _, err := io.ReadFull(r, prefix); err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading length prefix: %w", err)
+ }
+ payloadLen := binary.BigEndian.Uint32(prefix)
+
+ // Read the payload.
+ payload := make([]byte, payloadLen)
+ if _, err := io.ReadFull(r, payload); err != nil {
Review Comment:
readFrame trusts the 4-byte length prefix and allocates a payload slice of that size with no upper bound. A malformed/hostile peer (or corrupted stream) can trigger huge allocations/OOM; consider enforcing a maximum frame size and rejecting/limiting reads when payloadLen exceeds it.
##########
go-sdk/pkg/execution/comms.go:
##########
@@ -0,0 +1,128 @@
+// 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 (
+ "fmt"
+ "io"
+ "log/slog"
+ "sync"
+ "sync/atomic"
+)
+
+// CoordinatorComm manages bidirectional communication with the Airflow supervisor
+// over a length-prefixed msgpack socket connection.
+type CoordinatorComm struct {
+ reader io.Reader
+ writer io.Writer
+ nextID atomic.Int32
+ logger *slog.Logger
+
+ wmu sync.Mutex // serialises writes
+ rmu sync.Mutex // serialises reads
+}
+
+// NewCoordinatorComm creates a new communication channel.
+func NewCoordinatorComm(reader io.Reader, writer io.Writer, logger *slog.Logger) *CoordinatorComm {
+ return &CoordinatorComm{
+ reader: reader,
+ writer: writer,
+ logger: logger,
+ }
+}
+
+// ReadMessage reads and decodes one frame from the comm socket.
+// It returns the raw IncomingFrame with decoded map bodies.
+func (c *CoordinatorComm) ReadMessage() (IncomingFrame, error) {
+ c.rmu.Lock()
+ defer c.rmu.Unlock()
+ frame, err := readFrame(c.reader)
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading frame: %w", err)
+ }
+ c.logger.Debug("Received frame", "id", frame.ID)
+ return frame, nil
+}
+
+// SendRequest sends a request frame (2-element: [id, body]) to the supervisor.
+func (c *CoordinatorComm) SendRequest(id int, body map[string]any) error {
+ payload, err := encodeRequest(id, body)
+ if err != nil {
+ return fmt.Errorf("encoding request: %w", err)
+ }
+ c.logger.Debug("Sending request", "id", id)
+ c.wmu.Lock()
+ defer c.wmu.Unlock()
+ return writeFrame(c.writer, payload)
+}
+
+// Communicate sends a request and waits for the corresponding response.
+// This is a synchronous request-response: the caller sends a request and blocks
+// until the next frame arrives. The protocol is single-threaded on the comm socket.
+//
+// If the response contains an error element, it is returned as an ApiError.
+// Otherwise, the response body map is returned.
+func (c *CoordinatorComm) Communicate(body map[string]any) (map[string]any, error) {
+ id := int(c.nextID.Add(1) - 1)
+
+ if err := c.SendRequest(id, body); err != nil {
+ return nil, err
Review Comment:
CoordinatorComm.Communicate assumes the next frame read is the response for the request it just sent, but it never verifies frame.ID matches the generated id. This will break if calls happen concurrently or if the supervisor can send frames out-of-order; either serialize the full send+read with a single mutex (or implement a per-id dispatcher) and validate the response id before returning.
##########
go-sdk/pkg/execution/task_runner.go:
##########
@@ -0,0 +1,117 @@
+// 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"
+ "fmt"
+ "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, details)
+
+ // taskFunction.sendXcom reads the workload from context to get the task
+ // instance ids; populate it the same shape the gRPC path uses.
+ tiUUID, _ := uuid.Parse(details.TI.ID)
Review Comment:
uuid.Parse(details.TI.ID) ignores the returned error, which means an invalid TI id from the supervisor will silently become the zero UUID and be used in the workload context (impacting XCom pushes/heartbeats/etc). Please handle the parse error and return an appropriate terminal TaskState (or propagate a protocol-level error).
##########
go-sdk/pkg/execution/server.go:
##########
@@ -0,0 +1,155 @@
+// 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 selects between two
+// sub-protocols:
+//
+// - DagFileParseRequest: one-shot, returns DagFileParsingResult and exits.
+// - StartupDetails: multi-round task execution.
+//
+// See go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md.
+package execution
+
+import (
+ "fmt"
+ "log/slog"
+ "net"
+ "sync"
+
+ "github.com/apache/airflow/go-sdk/bundle/bundlev1"
+)
+
+// 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 (one-shot DAG parse or 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.
+ var commConn, logsConn net.Conn
+ var commErr, logsErr error
+ var wg sync.WaitGroup
+ wg.Add(2)
+ go func() {
+ defer wg.Done()
+ commConn, commErr = net.Dial("tcp", commAddr)
+ }()
Review Comment:
net.Dial has no timeout and can block indefinitely if the supervisor address is unreachable. Consider using net.Dialer with a reasonable Timeout (and/or context) for both comm/logs sockets so coordinator-mode startup fails fast rather than hanging.
##########
go-sdk/bundle/bundlev1/bundlev1server/server.go:
##########
@@ -52,41 +69,97 @@ 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.
+//
+// The mode is decided from CLI flags and process environment, so user code is
+// always one line:
//
-// 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
+// func main() { bundlev1server.Serve(&myBundle{}) }
+//
+// 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{}
for _, c := range opts {
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:
+ fmt.Fprintln(os.Stderr, "error: --comm and --logs must be supplied together")
+ flag.CommandLine.SetOutput(os.Stderr)
+ flag.Usage()
+ os.Exit(2)
Review Comment:
Serve() has an error-returning signature but calls os.Exit(2) on flag misuse, which makes it hard to test and prevents callers from handling the error. Consider returning an error (and letting main decide whether to print usage/exit) instead of exiting inside the library entrypoint.
##########
go-sdk/pkg/execution/messages.go:
##########
@@ -0,0 +1,407 @@
+// 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 (
+ "fmt"
+ "time"
+)
+
+// Inbound messages (Supervisor -> Runtime).
+
+// DagFileParseRequest is sent by the supervisor to request DAG parsing.
+type DagFileParseRequest struct {
+ File string
+ BundlePath string
+}
+
+func decodeDagFileParseRequest(m map[string]any) (*DagFileParseRequest, error) {
+ file, err := mapString(m, "file")
+ if err != nil {
+ return nil, err
+ }
+ bundlePath := mapStringOr(m, "bundle_path", "")
+ return &DagFileParseRequest{File: file, BundlePath: bundlePath}, nil
+}
+
+// TaskInstanceInfo holds task instance details from StartupDetails.
+type TaskInstanceInfo struct {
+ ID string
+ TaskID string
+ DagID string
+ RunID string
+ TryNumber int
+ DagVersionID string
+ MapIndex int
+ ContextCarrier map[string]any
+}
+
+func decodeTaskInstanceInfo(m map[string]any) (TaskInstanceInfo, error) {
+ if m == nil {
+ return TaskInstanceInfo{}, fmt.Errorf("nil task instance map")
+ }
+ id, err := mapString(m, "id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.id: %w", err)
+ }
+ taskID, err := mapString(m, "task_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.task_id: %w", err)
+ }
+ dagID, err := mapString(m, "dag_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.dag_id: %w", err)
+ }
+ runID, err := mapString(m, "run_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.run_id: %w", err)
+ }
+ tryNumber := mapIntOr(m, "try_number", 1)
+ dagVersionID := mapStringOr(m, "dag_version_id", "")
+ mapIndex := mapIntOr(m, "map_index", -1)
+ contextCarrier := mapMap(m, "context_carrier")
+
+ return TaskInstanceInfo{
+ ID: id,
+ TaskID: taskID,
+ DagID: dagID,
+ RunID: runID,
+ TryNumber: tryNumber,
+ DagVersionID: dagVersionID,
+ MapIndex: mapIndex,
+ ContextCarrier: contextCarrier,
+ }, nil
+}
+
+// BundleInfoMsg holds bundle identification from StartupDetails.
+type BundleInfoMsg struct {
+ Name string
+ Version string
+}
+
+func decodeBundleInfo(m map[string]any) BundleInfoMsg {
+ if m == nil {
+ return BundleInfoMsg{}
+ }
+ return BundleInfoMsg{
+ Name: mapStringOr(m, "name", ""),
+ Version: mapStringOr(m, "version", ""),
+ }
+}
+
+// TIRunContext holds the runtime context for a task instance.
+type TIRunContext struct {
+ LogicalDate *time.Time
+ DataIntervalStart *time.Time
+ DataIntervalEnd *time.Time
+}
+
+func decodeTIRunContext(m map[string]any) TIRunContext {
+ if m == nil {
+ return TIRunContext{}
+ }
+ ctx := TIRunContext{}
+ if t, err := asTime(m["logical_date"]); err == nil {
+ ctx.LogicalDate = &t
+ }
+ if t, err := asTime(m["data_interval_start"]); err == nil {
+ ctx.DataIntervalStart = &t
+ }
+ if t, err := asTime(m["data_interval_end"]); err == nil {
+ ctx.DataIntervalEnd = &t
+ }
+ return ctx
+}
+
+// StartupDetails is sent by the supervisor to initiate task execution.
+type StartupDetails struct {
+ TI TaskInstanceInfo
+ DagRelPath string
+ BundleInfo BundleInfoMsg
+ StartDate time.Time
+ TIContext TIRunContext
+ SentryIntegration string
+}
+
+func decodeStartupDetails(m map[string]any) (*StartupDetails, error) {
+ tiMap := mapMap(m, "ti")
+ ti, err := decodeTaskInstanceInfo(tiMap)
+ if err != nil {
+ return nil, fmt.Errorf("decoding ti: %w", err)
+ }
+
+ dagRelPath := mapStringOr(m, "dag_rel_path", "")
+ bundleInfo := decodeBundleInfo(mapMap(m, "bundle_info"))
+
+ var startDate time.Time
+ if t, err := asTime(m["start_date"]); err == nil {
+ startDate = t
Review Comment:
decodeStartupDetails ignores start_date parse errors and leaves StartDate as the zero time. If start_date is provided but malformed, this should likely be treated as a decoding error (or at least logged) to avoid silently running tasks with an incorrect start date.
##########
go-sdk/pkg/execution/messages.go:
##########
@@ -0,0 +1,407 @@
+// 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 (
+ "fmt"
+ "time"
+)
+
+// Inbound messages (Supervisor -> Runtime).
+
+// DagFileParseRequest is sent by the supervisor to request DAG parsing.
+type DagFileParseRequest struct {
+ File string
+ BundlePath string
+}
+
+func decodeDagFileParseRequest(m map[string]any) (*DagFileParseRequest, error) {
+ file, err := mapString(m, "file")
+ if err != nil {
+ return nil, err
+ }
+ bundlePath := mapStringOr(m, "bundle_path", "")
+ return &DagFileParseRequest{File: file, BundlePath: bundlePath}, nil
+}
+
+// TaskInstanceInfo holds task instance details from StartupDetails.
+type TaskInstanceInfo struct {
+ ID string
+ TaskID string
+ DagID string
+ RunID string
+ TryNumber int
+ DagVersionID string
+ MapIndex int
+ ContextCarrier map[string]any
+}
+
+func decodeTaskInstanceInfo(m map[string]any) (TaskInstanceInfo, error) {
+ if m == nil {
+ return TaskInstanceInfo{}, fmt.Errorf("nil task instance map")
+ }
+ id, err := mapString(m, "id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.id: %w", err)
+ }
+ taskID, err := mapString(m, "task_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.task_id: %w", err)
+ }
+ dagID, err := mapString(m, "dag_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.dag_id: %w", err)
+ }
+ runID, err := mapString(m, "run_id")
+ if err != nil {
+ return TaskInstanceInfo{}, fmt.Errorf("ti.run_id: %w", err)
+ }
+ tryNumber := mapIntOr(m, "try_number", 1)
+ dagVersionID := mapStringOr(m, "dag_version_id", "")
+ mapIndex := mapIntOr(m, "map_index", -1)
+ contextCarrier := mapMap(m, "context_carrier")
+
+ return TaskInstanceInfo{
+ ID: id,
+ TaskID: taskID,
+ DagID: dagID,
+ RunID: runID,
+ TryNumber: tryNumber,
+ DagVersionID: dagVersionID,
+ MapIndex: mapIndex,
+ ContextCarrier: contextCarrier,
+ }, nil
+}
+
+// BundleInfoMsg holds bundle identification from StartupDetails.
+type BundleInfoMsg struct {
+ Name string
+ Version string
+}
+
+func decodeBundleInfo(m map[string]any) BundleInfoMsg {
+ if m == nil {
+ return BundleInfoMsg{}
+ }
+ return BundleInfoMsg{
+ Name: mapStringOr(m, "name", ""),
+ Version: mapStringOr(m, "version", ""),
+ }
+}
+
+// TIRunContext holds the runtime context for a task instance.
+type TIRunContext struct {
+ LogicalDate *time.Time
+ DataIntervalStart *time.Time
+ DataIntervalEnd *time.Time
+}
+
+func decodeTIRunContext(m map[string]any) TIRunContext {
+ if m == nil {
+ return TIRunContext{}
+ }
+ ctx := TIRunContext{}
+ if t, err := asTime(m["logical_date"]); err == nil {
+ ctx.LogicalDate = &t
+ }
+ if t, err := asTime(m["data_interval_start"]); err == nil {
+ ctx.DataIntervalStart = &t
+ }
+ if t, err := asTime(m["data_interval_end"]); err == nil {
+ ctx.DataIntervalEnd = &t
+ }
+ return ctx
Review Comment:
decodeTIRunContext silently drops logical_date/data_interval_* values when asTime fails (wrong type/invalid RFC3339). For protocol messages it’s usually better to return an explicit decoding error when a field is present but malformed, so mismatches don’t get hidden as nil timestamps.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [For review only, Don't merge] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 closed pull request #67154: [For review only, Don't merge] Go SDK: coordinator-protocol execution server and task runner
URL: https://github.com/apache/airflow/pull/67154
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #67154:
URL: https://github.com/apache/airflow/pull/67154#discussion_r3285594567
##########
go-sdk/pkg/execution/server.go:
##########
@@ -0,0 +1,164 @@
+// 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 selects between two
+// sub-protocols:
+//
+// - DagFileParseRequest: one-shot, returns DagFileParsingResult and exits.
+// - StartupDetails: 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 (one-shot DAG parse or 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()
+
+ if commErr != nil {
+ return fmt.Errorf("connecting to comm socket %s: %w", commAddr, commErr)
+ }
+ defer commConn.Close()
+ if logsErr != nil {
+ return fmt.Errorf("connecting to logs socket %s: %w", logsAddr, logsErr)
+ }
+ defer logsConn.Close()
Review Comment:
When the comm dial fails, the logs dial may still have succeeded. Returning immediately without closing `logsConn` can leak an open TCP connection (and similarly, `commConn` can be non-nil even when `commErr` is set).
Close any non-nil conns on the error paths before returning.
##########
go-sdk/pkg/execution/frames.go:
##########
@@ -0,0 +1,249 @@
+// 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"
+ "encoding/binary"
+ "fmt"
+ "io"
+
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// maxFrameSize caps the payload length a single frame may declare. A
+// malformed length prefix from a corrupted stream (or hostile peer) would
+// otherwise let readFrame allocate up to 4 GiB before the read failed.
+// 64 MiB is far above any legitimate StartupDetails or XCom payload while
+// still preventing accidental OOM.
+const maxFrameSize = 64 * 1024 * 1024
+
+// IncomingFrame represents a decoded frame received from the comm socket.
+type IncomingFrame struct {
+ ID int
+ Body map[string]any
+ Err map[string]any // non-nil only for response frames (3-element arrays)
+}
+
+// encodeRequest encodes a request frame (2-element msgpack array: [id, body]).
+func encodeRequest(id int, body map[string]any) ([]byte, error) {
+ var buf bytes.Buffer
+ enc := msgpack.NewEncoder(&buf)
+ enc.UseCompactInts(true)
+
+ if err := enc.EncodeArrayLen(2); err != nil {
+ return nil, err
+ }
+ if err := enc.EncodeInt(int64(id)); err != nil {
+ return nil, err
+ }
+ if err := enc.Encode(body); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+}
+
+// writeFrame writes a length-prefixed msgpack payload to the writer.
+// Format: [4-byte big-endian length][payload bytes]
+func writeFrame(w io.Writer, payload []byte) error {
+ prefix := make([]byte, 4)
+ binary.BigEndian.PutUint32(prefix, uint32(len(payload)))
+ if _, err := w.Write(prefix); err != nil {
+ return fmt.Errorf("writing length prefix: %w", err)
+ }
+ if _, err := w.Write(payload); err != nil {
+ return fmt.Errorf("writing payload: %w", err)
+ }
+ return nil
Review Comment:
`writeFrame` assumes `io.Writer.Write` writes the full buffer in one call. The `io.Writer` contract allows short writes with `n < len(p)` and `err == nil`, which would corrupt framing on the wire.
Check `n` and surface `io.ErrShortWrite` (or loop until the full buffer is written).
##########
go-sdk/pkg/execution/frames.go:
##########
@@ -0,0 +1,249 @@
+// 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"
+ "encoding/binary"
+ "fmt"
+ "io"
+
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// maxFrameSize caps the payload length a single frame may declare. A
+// malformed length prefix from a corrupted stream (or hostile peer) would
+// otherwise let readFrame allocate up to 4 GiB before the read failed.
+// 64 MiB is far above any legitimate StartupDetails or XCom payload while
+// still preventing accidental OOM.
+const maxFrameSize = 64 * 1024 * 1024
+
+// IncomingFrame represents a decoded frame received from the comm socket.
+type IncomingFrame struct {
+ ID int
+ Body map[string]any
+ Err map[string]any // non-nil only for response frames (3-element arrays)
+}
+
+// encodeRequest encodes a request frame (2-element msgpack array: [id, body]).
+func encodeRequest(id int, body map[string]any) ([]byte, error) {
+ var buf bytes.Buffer
+ enc := msgpack.NewEncoder(&buf)
+ enc.UseCompactInts(true)
+
+ if err := enc.EncodeArrayLen(2); err != nil {
+ return nil, err
+ }
+ if err := enc.EncodeInt(int64(id)); err != nil {
+ return nil, err
+ }
+ if err := enc.Encode(body); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+}
+
+// writeFrame writes a length-prefixed msgpack payload to the writer.
+// Format: [4-byte big-endian length][payload bytes]
+func writeFrame(w io.Writer, payload []byte) error {
+ prefix := make([]byte, 4)
+ binary.BigEndian.PutUint32(prefix, uint32(len(payload)))
+ if _, err := w.Write(prefix); err != nil {
+ return fmt.Errorf("writing length prefix: %w", err)
+ }
+ if _, err := w.Write(payload); err != nil {
+ return fmt.Errorf("writing payload: %w", err)
+ }
+ return nil
+}
+
+// readFrame reads one length-prefixed msgpack frame from the reader and decodes it.
+func readFrame(r io.Reader) (IncomingFrame, error) {
+ // Read 4-byte big-endian length prefix.
+ prefix := make([]byte, 4)
+ if _, err := io.ReadFull(r, prefix); err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading length prefix: %w", err)
+ }
+ payloadLen := binary.BigEndian.Uint32(prefix)
+ if payloadLen > maxFrameSize {
+ return IncomingFrame{}, fmt.Errorf(
+ "frame payload length %d exceeds max %d",
+ payloadLen,
+ maxFrameSize,
+ )
+ }
+
+ // Read the payload.
+ payload := make([]byte, payloadLen)
+ if _, err := io.ReadFull(r, payload); err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading payload (%d bytes): %w", payloadLen, err)
+ }
Review Comment:
`payloadLen` is a `uint32`, but `make([]byte, payloadLen)` requires an `int` length. This currently won’t compile.
Convert to `int(payloadLen)` after the max-size guard (so the conversion is safe) before allocating/reading the payload.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Go SDK: coordinator-protocol execution server and task runner [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #67154:
URL: https://github.com/apache/airflow/pull/67154#discussion_r3280509008
##########
go-sdk/sdk/sdk.go:
##########
@@ -43,6 +52,12 @@ type VariableClient interface {
// // Other errors here, such as http network timeouts etc.
// }
GetVariable(ctx context.Context, key string) (string, error)
+
+ // UnmarshalJSONVariable fetches a variable and json.Unmarshal's its value
+ // into pointer. Use this when the variable was stored as a JSON object,
+ // array, or number; for plain string variables call GetVariable directly.
Review Comment:
Docstring grammar: "json.Unmarshal's" is unidiomatic/incorrect (reads like a possessive). Consider rephrasing to "json.Unmarshal" (e.g. "unmarshals") to keep GoDoc clear.
##########
go-sdk/pkg/execution/task_runner.go:
##########
@@ -0,0 +1,126 @@
+// 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"
+ "fmt"
+ "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, details)
+
+ // 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.Error("Task failed", "error", fmt.Sprintf("%v", err))
Review Comment:
Logging the error via fmt.Sprintf("%v", err) loses the structured error value; slog can take the error directly. Pass the error as the value (and ideally use ErrorContext with ctx) so downstream handlers can format/redact consistently.
##########
go-sdk/pkg/execution/client.go:
##########
@@ -0,0 +1,201 @@
+// 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"
+ "encoding/json"
+ "fmt"
+
+ "github.com/apache/airflow/go-sdk/pkg/api"
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+// CoordinatorClient implements sdk.Client by communicating with the Airflow supervisor
+// over the comm socket using msgpack-framed IPC instead of HTTP.
+type CoordinatorClient struct {
+ comm *CoordinatorComm
+ details *StartupDetails
+}
+
+var _ sdk.Client = (*CoordinatorClient)(nil)
+
+// NewCoordinatorClient creates a new client backed by the comm socket.
+func NewCoordinatorClient(comm *CoordinatorComm, details *StartupDetails) *CoordinatorClient {
+ return &CoordinatorClient{
+ comm: comm,
+ details: details,
+ }
+}
+
+// GetVariable requests a variable value from the supervisor.
+func (c *CoordinatorClient) GetVariable(_ context.Context, key string) (string, error) {
+ resp, err := c.comm.Communicate(GetVariableMsg{Key: key}.toMap())
Review Comment:
CoordinatorClient methods currently ignore the provided context (e.g. GetVariable takes `_ context.Context`). If the task context is cancelled or has a deadline, SDK calls can still block indefinitely waiting on the comm socket. Thread ctx through to CoordinatorComm (or at least check ctx.Done before/while waiting for the response) so cancellations propagate cleanly.
##########
go-sdk/pkg/execution/comms.go:
##########
@@ -0,0 +1,226 @@
+// 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 (
+ "errors"
+ "fmt"
+ "io"
+ "log/slog"
+ "sync"
+ "sync/atomic"
+)
+
+// CoordinatorComm manages bidirectional communication with the Airflow
+// supervisor over a length-prefixed msgpack socket connection.
+//
+// Reads are multiplexed by frame ID. A single background reader goroutine,
+// lazily started on the first Communicate call, consumes inbound frames and
+// dispatches each one to the Communicate caller whose request used that ID.
+// This lets multiple goroutines call Communicate concurrently without
+// serialising the full send-then-read round trip behind a single mutex, and
+// guarantees the response a caller receives matches the request it sent.
+//
+// The supervisor's initial frame (DagFileParseRequest or StartupDetails)
+// arrives unsolicited, before any client request is in flight, and is read
+// synchronously via ReadMessage; ReadMessage must not be called after the
+// dispatcher has been started.
+type CoordinatorComm struct {
+ reader io.Reader
+ writer io.Writer
+ nextID atomic.Int64
+ logger *slog.Logger
+
+ wmu sync.Mutex // serialises writes
+
+ // Multiplexer state. mu protects every field below.
+ mu sync.Mutex
+ pending map[int]chan frameResult
+ started bool
+ readErr error
+}
+
+// frameResult is the value the dispatcher delivers to a Communicate caller.
+type frameResult struct {
+ frame IncomingFrame
+ err error
+}
+
+// ErrDispatcherClosed is wrapped into the error Communicate returns once the
+// background reader goroutine has exited — typically because the supervisor
+// closed the comm socket.
+var ErrDispatcherClosed = errors.New("coordinator comm: dispatcher closed")
+
+// NewCoordinatorComm creates a new communication channel.
+func NewCoordinatorComm(reader io.Reader, writer io.Writer, logger *slog.Logger) *CoordinatorComm {
+ return &CoordinatorComm{
+ reader: reader,
+ writer: writer,
+ logger: logger,
+ pending: make(map[int]chan frameResult),
+ }
+}
+
+// ReadMessage reads and decodes one frame directly from the comm socket.
+// It is used to read the supervisor's initial frame before any request/response
+// traffic begins. Calling it after the dispatcher has started would race the
+// reader goroutine for input bytes, so it returns an error in that case.
+func (c *CoordinatorComm) ReadMessage() (IncomingFrame, error) {
+ c.mu.Lock()
+ started := c.started
+ c.mu.Unlock()
+ if started {
+ return IncomingFrame{}, errors.New(
+ "coordinator comm: ReadMessage cannot be used after the dispatcher has started",
+ )
+ }
+
+ frame, err := readFrame(c.reader)
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading frame: %w", err)
+ }
+ c.logger.Debug("Received frame", "id", frame.ID)
+ return frame, nil
+}
+
+// SendRequest writes a request frame (2-element [id, body]) to the supervisor.
+// Concurrent calls are serialised so frames are never interleaved on the wire.
+func (c *CoordinatorComm) SendRequest(id int, body map[string]any) error {
+ payload, err := encodeRequest(id, body)
+ if err != nil {
+ return fmt.Errorf("encoding request: %w", err)
+ }
+ c.logger.Debug("Sending request", "id", id)
+ c.wmu.Lock()
+ defer c.wmu.Unlock()
+ return writeFrame(c.writer, payload)
+}
+
+// Communicate sends a request and blocks until the supervisor's response with
+// the matching frame ID is delivered by the dispatcher. Safe to call
+// concurrently from multiple goroutines.
+//
+// If the response carries an error (either as the third element of a 3-tuple
+// frame or as a body whose "type" is "ErrorResponse") it is returned as an
+// *ApiError. If the dispatcher's read loop has terminated, the underlying read
+// error is returned wrapped in ErrDispatcherClosed.
+func (c *CoordinatorComm) Communicate(body map[string]any) (map[string]any, error) {
+ id := int(c.nextID.Add(1) - 1)
+ ch := make(chan frameResult, 1)
+
+ // Register the waiter under the same lock the dispatcher uses, and before
+ // sending the request, so the dispatcher can never deliver the response
+ // (or a terminal read error) before this caller is ready to receive it.
+ c.mu.Lock()
+ if !c.started {
+ c.started = true
+ go c.readLoop()
+ }
+ if c.readErr != nil {
+ err := c.readErr
+ c.mu.Unlock()
+ return nil, fmt.Errorf("%w: %w", ErrDispatcherClosed, err)
+ }
+ c.pending[id] = ch
+ c.mu.Unlock()
+
+ if err := c.SendRequest(id, body); err != nil {
+ c.mu.Lock()
+ delete(c.pending, id)
+ c.mu.Unlock()
+ return nil, err
+ }
+
+ result := <-ch
+ if result.err != nil {
+ return nil, fmt.Errorf("%w: %w", ErrDispatcherClosed, result.err)
Review Comment:
CoordinatorComm.Communicate blocks on `<-ch` with no way to respect cancellation/timeouts. Since higher-level SDK methods accept a context, consider adding `ctx context.Context` to Communicate and selecting on `ctx.Done()` to avoid permanently stuck goroutines when the supervisor is slow/unresponsive.
##########
go-sdk/pkg/execution/frames.go:
##########
@@ -0,0 +1,243 @@
+// 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"
+ "encoding/binary"
+ "fmt"
+ "io"
+
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// maxFrameSize caps the payload length a single frame may declare. A
+// malformed length prefix from a corrupted stream (or hostile peer) would
+// otherwise let readFrame allocate up to 4 GiB before the read failed.
+// 64 MiB is far above any legitimate StartupDetails or XCom payload while
+// still preventing accidental OOM.
+const maxFrameSize = 64 * 1024 * 1024
+
+// IncomingFrame represents a decoded frame received from the comm socket.
+type IncomingFrame struct {
+ ID int
+ Body map[string]any
+ Err map[string]any // non-nil only for response frames (3-element arrays)
+}
+
+// encodeRequest encodes a request frame (2-element msgpack array: [id, body]).
+func encodeRequest(id int, body map[string]any) ([]byte, error) {
+ var buf bytes.Buffer
+ enc := msgpack.NewEncoder(&buf)
+ enc.UseCompactInts(true)
+
+ if err := enc.EncodeArrayLen(2); err != nil {
+ return nil, err
+ }
+ if err := enc.EncodeInt(int64(id)); err != nil {
+ return nil, err
+ }
+ if err := enc.Encode(body); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+}
+
+// writeFrame writes a length-prefixed msgpack payload to the writer.
+// Format: [4-byte big-endian length][payload bytes]
+func writeFrame(w io.Writer, payload []byte) error {
+ prefix := make([]byte, 4)
+ binary.BigEndian.PutUint32(prefix, uint32(len(payload)))
+ if _, err := w.Write(prefix); err != nil {
+ return fmt.Errorf("writing length prefix: %w", err)
+ }
+ if _, err := w.Write(payload); err != nil {
+ return fmt.Errorf("writing payload: %w", err)
+ }
+ return nil
+}
+
+// readFrame reads one length-prefixed msgpack frame from the reader and decodes it.
+func readFrame(r io.Reader) (IncomingFrame, error) {
+ // Read 4-byte big-endian length prefix.
+ prefix := make([]byte, 4)
+ if _, err := io.ReadFull(r, prefix); err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading length prefix: %w", err)
+ }
+ payloadLen := binary.BigEndian.Uint32(prefix)
+ if payloadLen > maxFrameSize {
+ return IncomingFrame{}, fmt.Errorf(
+ "frame payload length %d exceeds max %d",
+ payloadLen,
+ maxFrameSize,
+ )
+ }
+
+ // Read the payload.
+ payload := make([]byte, payloadLen)
+ if _, err := io.ReadFull(r, payload); err != nil {
+ return IncomingFrame{}, fmt.Errorf("reading payload (%d bytes): %w", payloadLen, err)
+ }
+
+ return decodeFrame(payload)
+}
+
+// decodeFrame decodes a msgpack payload into an IncomingFrame.
+func decodeFrame(data []byte) (IncomingFrame, error) {
+ dec := msgpack.NewDecoder(bytes.NewReader(data))
+
+ arrLen, err := dec.DecodeArrayLen()
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("decoding array header: %w", err)
+ }
+ if arrLen < 2 {
+ return IncomingFrame{}, fmt.Errorf("unexpected frame arity %d, need at least 2", arrLen)
+ }
+
+ id64, err := dec.DecodeInt64()
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("decoding frame id: %w", err)
+ }
+
+ // Decode the body element.
+ bodyRaw, err := dec.DecodeInterface()
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("decoding body: %w", err)
+ }
+ body, _ := toStringMap(bodyRaw)
+
+ // For response frames (3-element), decode the error element.
+ var errMap map[string]any
+ if arrLen >= 3 {
+ errRaw, err := dec.DecodeInterface()
+ if err != nil {
+ return IncomingFrame{}, fmt.Errorf("decoding error element: %w", err)
+ }
+ errMap, _ = toStringMap(errRaw)
+ }
Review Comment:
decodeFrame currently ignores failed map conversions (`body, _ := toStringMap(bodyRaw)` / `errMap, _ = toStringMap(errRaw)`), silently turning a non-map payload into nil. This can mask protocol bugs and lead to confusing downstream errors. Consider returning a decode error when bodyRaw/errRaw is non-nil but not a map.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org