uranusjr commented on code in PR #69079: URL: https://github.com/apache/airflow/pull/69079#discussion_r3489981675
########## ts-sdk/src/coordinator/comm-channel.ts: ########## @@ -0,0 +1,265 @@ +/*! + * 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. + */ + +// Comm socket client — length-prefixed msgpack frames over TCP. +// Mirrors the Airflow supervisor's comm socket protocol. +// +// The channel is the sole reader on the socket. The task sends +// requests and awaits id-correlated replies. The supervisor's only +// unprompted frame is the greeting (StartupDetails / +// DagFileParseRequest), pre-caught into the `greeting` promise that +// `connect()` awaits — the protocol sends nothing else +// supervisor-initiated (comms.py: "No messages are sent to task +// process except in response to a request"). + +import type { Socket } from "node:net"; +import { encodeRequest, encodeResponse, FrameReader, type Frame } from "./frames.js"; +import { connectTcp } from "./tcp-connect.js"; +import { Deferred } from "./deferred.js"; +import type { LogChannel } from "./log-channel.js"; + +/** What `CommChannel.connect` resolves to: the live channel plus the + * supervisor's first frame (StartupDetails / DagFileParseRequest), + * already in hand so the caller never has to manage a "frame arrived + * with no consumer" window. */ +export interface CommConnection { + channel: CommChannel; + firstFrame: Frame; +} + +export interface SendResponseOptions { + timeoutMs?: number; +} + +export class CommChannel { + private readonly sock: Socket; + private readonly reader = new FrameReader(); + private readonly logs: LogChannel | null; + private nextId = 0; + private pendingReplies = new Map<number, (frame: Frame) => void>(); + private closed = false; + private closeError: Error | null = null; + + // The greeting (first supervisor-initiated frame). The promise is + // its own buffer: arriving before `connect()` awaits is fine — it + // stays settled with the value, so there is no race to handle. + private readonly greeting = new Deferred<Frame>(); + + private constructor(sock: Socket, logs: LogChannel | null) { + this.sock = sock; + this.logs = logs; + // A `new Socket()` (from `connectTcp`) starts paused: it buffers + // inbound bytes and emits no `data` until a listener attaches + // and flips it to flowing. Attaching synchronously here — same + // tick as construction, before the event loop can deliver a + // read, and as the only reader — loses nothing, double-reads + // nothing. + sock.on("data", (chunk) => this.handleData(chunk)); + sock.on("close", () => this.handleClose(null)); + sock.on("error", (err) => this.handleClose(err)); + } + + /** Connect and wait for the supervisor's greeting; rejects if the + * socket dies before it arrives. */ + static async connect(addr: string, logs: LogChannel | null = null): Promise<CommConnection> { + const sock = await connectTcp(addr); + const channel = new CommChannel(sock, logs); + const firstFrame = await channel.greeting.promise; + return { channel, firstFrame }; + } + + /** Send a request to the supervisor and await its matching response. */ + async request(body: unknown): Promise<Frame> { + const id = this.nextId++; + const type = describeFrameType(body); + this.logs?.debug("Sending request", { id, type }); + return new Promise<Frame>((resolve, reject) => { + if (this.closed) { + reject(this.closeError ?? new Error("Comm channel closed")); + return; + } + this.pendingReplies.set(id, (frame) => { + this.logs?.debug("Response received", { + id, + request_type: type, + response_type: describeFrameType(frame.body), + error: frame.error ?? null, + }); + resolve(frame); + }); Review Comment: Not sure if I missed anything. This has the potential to be stuck forever (until the process is killed) if the frame it expects got dropped or corrupt for whatever reason. This should probably add a timeout or something. ########## ts-sdk/src/coordinator/comm-channel.ts: ########## @@ -0,0 +1,265 @@ +/*! + * 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. + */ + +// Comm socket client — length-prefixed msgpack frames over TCP. +// Mirrors the Airflow supervisor's comm socket protocol. +// +// The channel is the sole reader on the socket. The task sends +// requests and awaits id-correlated replies. The supervisor's only +// unprompted frame is the greeting (StartupDetails / +// DagFileParseRequest), pre-caught into the `greeting` promise that +// `connect()` awaits — the protocol sends nothing else +// supervisor-initiated (comms.py: "No messages are sent to task +// process except in response to a request"). + +import type { Socket } from "node:net"; +import { encodeRequest, encodeResponse, FrameReader, type Frame } from "./frames.js"; +import { connectTcp } from "./tcp-connect.js"; +import { Deferred } from "./deferred.js"; +import type { LogChannel } from "./log-channel.js"; + +/** What `CommChannel.connect` resolves to: the live channel plus the + * supervisor's first frame (StartupDetails / DagFileParseRequest), + * already in hand so the caller never has to manage a "frame arrived + * with no consumer" window. */ +export interface CommConnection { + channel: CommChannel; + firstFrame: Frame; +} + +export interface SendResponseOptions { + timeoutMs?: number; +} + +export class CommChannel { + private readonly sock: Socket; + private readonly reader = new FrameReader(); + private readonly logs: LogChannel | null; + private nextId = 0; + private pendingReplies = new Map<number, (frame: Frame) => void>(); + private closed = false; + private closeError: Error | null = null; + + // The greeting (first supervisor-initiated frame). The promise is + // its own buffer: arriving before `connect()` awaits is fine — it + // stays settled with the value, so there is no race to handle. + private readonly greeting = new Deferred<Frame>(); + + private constructor(sock: Socket, logs: LogChannel | null) { + this.sock = sock; + this.logs = logs; + // A `new Socket()` (from `connectTcp`) starts paused: it buffers + // inbound bytes and emits no `data` until a listener attaches + // and flips it to flowing. Attaching synchronously here — same + // tick as construction, before the event loop can deliver a + // read, and as the only reader — loses nothing, double-reads + // nothing. + sock.on("data", (chunk) => this.handleData(chunk)); + sock.on("close", () => this.handleClose(null)); + sock.on("error", (err) => this.handleClose(err)); + } + + /** Connect and wait for the supervisor's greeting; rejects if the + * socket dies before it arrives. */ + static async connect(addr: string, logs: LogChannel | null = null): Promise<CommConnection> { + const sock = await connectTcp(addr); + const channel = new CommChannel(sock, logs); + const firstFrame = await channel.greeting.promise; + return { channel, firstFrame }; + } + + /** Send a request to the supervisor and await its matching response. */ + async request(body: unknown): Promise<Frame> { + const id = this.nextId++; + const type = describeFrameType(body); + this.logs?.debug("Sending request", { id, type }); + return new Promise<Frame>((resolve, reject) => { + if (this.closed) { + reject(this.closeError ?? new Error("Comm channel closed")); + return; + } + this.pendingReplies.set(id, (frame) => { + this.logs?.debug("Response received", { + id, + request_type: type, + response_type: describeFrameType(frame.body), + error: frame.error ?? null, + }); + resolve(frame); + }); Review Comment: Not sure if I missed anything. This seems to have the potential to be stuck forever (until the process is killed) if the frame it expects got dropped or corrupt for whatever reason. This should probably add a timeout or something. -- 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]
