kaxil commented on code in PR #67317: URL: https://github.com/apache/airflow/pull/67317#discussion_r3325519524
########## go-sdk/pkg/execution/logger.go: ########## @@ -0,0 +1,215 @@ +// 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" + "io" + "log/slog" + "strings" + "sync" + "time" +) + +// SocketLogHandler is an slog.Handler that streams structured JSON log lines +// to the logs TCP socket. Each log entry is a single JSON object followed by +// a newline, matching the Airflow log streaming format. +// +// Key mapping: +// - "event" for the log message (not "msg") +// - "level" in lowercase (not "INFO"/"ERROR") +// - "timestamp" in RFC3339Nano format (not "time") +// - Additional attributes are included as top-level fields +// +// Groups are encoded as dotted key prefixes on a flat JSON object +// (`{"grp.key": "val"}`), not as nested objects. The Airflow supervisor's +// log-streaming format consumes flat top-level fields, so emitting nested +// objects here would not be parsed correctly. Do not change this without +// updating the supervisor side in lockstep. +type SocketLogHandler struct { + shared *socketLogHandlerShared + level slog.Level + // attrs is the list of attributes accumulated via WithAttrs. Each entry's + // key has already been qualified with whatever groups were active at the + // WithAttrs call site, so a later WithGroup does NOT retroactively prefix + // them — matching the slog.Handler contract that groups apply only to + // subsequently-added attributes. + attrs []prefixedAttr + groups []string +} + +// prefixedAttr is an attribute whose key has been pre-qualified with the +// dotted prefix of the groups that were active when it was added via +// WithAttrs. The slog.Value is kept unresolved so LogValuer attributes are +// still evaluated lazily at Handle time. +type prefixedAttr struct { + key string + value slog.Value +} + +// socketLogHandlerShared holds the writer and buffer that must remain shared +// across WithAttrs / WithGroup clones; otherwise the sync.Mutex would be +// copied (which the runtime detector flags as a bug). +type socketLogHandlerShared struct { + mu sync.Mutex + writer io.Writer + buf [][]byte + connected bool +} + +var _ slog.Handler = (*SocketLogHandler)(nil) + +// NewSocketLogHandler creates a new handler. If writer is nil, messages are +// buffered until Connect() is called. +func NewSocketLogHandler(writer io.Writer, level slog.Level) *SocketLogHandler { + shared := &socketLogHandlerShared{} + if writer != nil { + shared.writer = writer + shared.connected = true + } + return &SocketLogHandler{ + shared: shared, + level: level, + } +} + +// Connect sets the writer and flushes any buffered log messages. +func (h *SocketLogHandler) Connect(w io.Writer) { + h.shared.mu.Lock() + defer h.shared.mu.Unlock() + + h.shared.writer = w + h.shared.connected = true + + for _, line := range h.shared.buf { + _, _ = w.Write(line) + } + h.shared.buf = nil +} + +func (h *SocketLogHandler) Enabled(_ context.Context, level slog.Level) bool { + return level >= h.level +} + +func (h *SocketLogHandler) Handle(_ context.Context, r slog.Record) error { + entry := make(map[string]any) + + // Set standard fields. + entry["event"] = r.Message + entry["level"] = strings.ToLower(r.Level.String()) + if !r.Time.IsZero() { + entry["timestamp"] = r.Time.Format(time.RFC3339Nano) + } + + // Apply pre-configured attrs. Keys are already qualified with the groups + // active at the WithAttrs call site, so the current h.groups is NOT + // applied here — only to record-level attrs below. + for _, a := range h.attrs { + entry[a.key] = resolveAttrValue(a.value) + } + + // Apply record attrs. These were added at Handle time, so they pick up + // the currently-active group prefix. + r.Attrs(func(a slog.Attr) bool { + entry[h.prefixedKey(a.Key)] = resolveAttrValue(a.Value) Review Comment: This flattens `WithGroup` correctly, but an inline group attr gets dropped. `logger.Info("msg", slog.Group("req", slog.String("method", "GET")))` (and `runner.go:274` does exactly this with `slog.Group("request", ...)` on the task-start error path) produces `"req":[{"Key":"method","Value":{}}]` here, versus `"req":{"method":"GET"}` from the stdlib JSON handler. `resolveAttrValue` calls `Resolve().Any()` on a `KindGroup` value, which returns `[]slog.Attr` whose unexported `slog.Value` marshals to `{}`, so the values are lost. Since the task logger is handed to user task code, any task that logs with `slog.Group` loses data once this is wired up. Could `Handle` check `a.Value.Kind() == slog.KindGroup` and recurse, expanding the group into prefixed keys (`req.method`) the way `WithGroup` does? ########## go-sdk/pkg/execution/comms.go: ########## @@ -0,0 +1,273 @@ +// 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" + "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 StartupDetails frame 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[int64]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[int64]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 int64, 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, ctx is cancelled, or +// its deadline expires. 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( + ctx context.Context, + body map[string]any, +) (map[string]any, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + id := 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() + + // Run the send in a goroutine so ctx cancellation can interrupt a blocked + // write. We deliberately do not touch the underlying connection on cancel + // (no SetWriteDeadline, no Close): manipulating the stream mid-write would + // either poison future writes with a stale deadline or risk a partial + // length-prefixed frame on the wire. Instead, we let the send goroutine + // run to completion in the background. If it eventually succeeds the + // supervisor's response has no waiter and is discarded by the dispatcher; + // if it fails the dispatcher will surface the error to other callers. + sendErr := make(chan error, 1) + go func() { + sendErr <- c.SendRequest(id, body) + }() + + select { + case err := <-sendErr: + if err != nil { + c.mu.Lock() + delete(c.pending, id) + c.mu.Unlock() + return nil, err + } + case <-ctx.Done(): + c.mu.Lock() + delete(c.pending, id) + c.mu.Unlock() + return nil, ctx.Err() + } + + var result frameResult + select { + case result = <-ch: + case <-ctx.Done(): + // Drop the waiter so a late supervisor reply is logged-and-discarded + // by the dispatcher rather than piling up in c.pending. The dispatcher + // may have already delivered the response into the buffered channel + // between the select and this delete; that delivery is harmless (the + // channel is buffered to 1 and the caller has already given up). + c.mu.Lock() + delete(c.pending, id) + c.mu.Unlock() + return nil, ctx.Err() + } + if result.err != nil { + return nil, fmt.Errorf("%w: %w", ErrDispatcherClosed, result.err) + } + frame := result.frame + + // An error reply can arrive in two shapes: the third element of a 3-tuple + // response frame, or the body of a 2-tuple frame whose "type" is + // "ErrorResponse". Pick whichever the supervisor used and decode once. + if errMap := errMapFromFrame(frame); errMap != nil { + errResp := decodeErrorResponse(errMap) + return nil, &ApiError{ + Err: errResp.Error, + Detail: errResp.Detail, + } + } + + return frame.Body, nil +} + +// errMapFromFrame returns the error-shaped map carried by a response frame, or +// nil if the frame is not an error reply. The supervisor may surface an error +// either as the dedicated third element of a 3-tuple frame (frame.Err) or as +// the body of a 2-tuple frame whose "type" is "ErrorResponse". +func errMapFromFrame(f IncomingFrame) map[string]any { + if f.Err != nil { + return f.Err + } + if f.Body != nil { + if typ, _ := f.Body["type"].(string); typ == "ErrorResponse" { + return f.Body + } + } + return nil +} + +// readLoop is the dispatcher: it reads frames from the comm socket and routes +// each one to the channel registered by the matching Communicate caller. When +// readFrame returns an error (typically io.EOF on supervisor shutdown), it +// fans the error out to every pending waiter and exits; any subsequent +// Communicate call sees the error via readErr and returns immediately. +func (c *CoordinatorComm) readLoop() { + for { + frame, err := readFrame(c.reader) + if err != nil { + c.mu.Lock() + c.readErr = err + pending := c.pending + c.pending = nil + c.mu.Unlock() + for _, ch := range pending { + ch <- frameResult{err: err} + } + return + } + c.logger.Debug("Received frame", "id", frame.ID) + + c.mu.Lock() + ch, ok := c.pending[frame.ID] + if ok { + delete(c.pending, frame.ID) + } + c.mu.Unlock() + if !ok { + c.logger.Warn("Discarding frame with no matching waiter", "id", frame.ID) Review Comment: This fires at Warn for every reply that arrives after its caller gave up. The cancellation path in `Communicate` deliberately drops the waiter and lets a late reply land here, so a workload with request deadlines will log this on the normal timeout path, not just on genuine protocol bugs. Debug seems closer to the actual severity. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
