ryerraguntla commented on code in PR #3519: URL: https://github.com/apache/iggy/pull/3519#discussion_r3608225946
########## gateways/kafka/src/server.rs: ########## @@ -0,0 +1,394 @@ +// 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. + +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use bytes::{BufMut, BytesMut}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::broadcast; +use tokio::time::{timeout, timeout_at}; +use tokio_util::task::TaskTracker; +use tracing::{debug, error, info, warn}; + +use crate::error::{KafkaProtocolError, Result}; +use crate::protocol::api::{ + BrokerAdvertise, DEFAULT_KAFKA_PORT, ERROR_INVALID_REQUEST, encode_error_only_response, + handle_request, +}; +use crate::protocol::codec::Decoder; +use crate::protocol::header::{ + RequestHeader, ResponseHeader, request_header_version, response_header_version, +}; +use std::io; + +const READ_CHUNK: usize = 65536; + +#[derive(Debug, Clone)] +pub struct ServerConfig { + pub bind_addr: String, + /// Hostname or IP advertised in Metadata (`KAFKA_ADVERTISED_HOST`). Required when `bind_addr` + /// uses a wildcard address (`0.0.0.0` / `::`). + pub advertised_host: Option<String>, + /// Port advertised in Metadata (`KAFKA_ADVERTISED_PORT`). Defaults to the bind port. + pub advertised_port: Option<u16>, + pub max_frame_size: usize, + pub read_timeout: Duration, + pub write_timeout: Duration, +} + +impl Default for ServerConfig { + fn default() -> Self { + Self { + bind_addr: format!("127.0.0.1:{DEFAULT_KAFKA_PORT}"), + advertised_host: None, + advertised_port: None, + max_frame_size: 8 * 1024 * 1024, + read_timeout: Duration::from_secs(15), + write_timeout: Duration::from_secs(10), + } + } +} + +impl BrokerAdvertise { + /// Resolve the broker endpoint advertised in Metadata. + /// + /// `local_addr` is the address the listener is actually bound to (from `listener.local_addr()`). + /// + /// # Errors + /// + /// Returns `InvalidConfig` when `advertised_host` is empty or the listener binds to a wildcard + /// without an explicit advertised host. + pub fn from_server_config(config: &ServerConfig, local_addr: SocketAddr) -> Result<Self> { + let port = config + .advertised_port + .map_or_else(|| i32::from(local_addr.port()), i32::from); + + let host = if let Some(ref advertised) = config.advertised_host { + let trimmed = advertised.trim(); + if trimmed.is_empty() { + return Err(KafkaProtocolError::InvalidConfig( + "KAFKA_ADVERTISED_HOST must not be empty".into(), + )); + } + if trimmed.len() > i16::MAX as usize { + return Err(KafkaProtocolError::InvalidConfig( + "KAFKA_ADVERTISED_HOST exceeds Kafka nullable string limit (32767 bytes)" + .into(), + )); + } + trimmed.to_string() + } else if local_addr.ip().is_unspecified() { + return Err(KafkaProtocolError::InvalidConfig( + "binding to a wildcard address (0.0.0.0 or ::) requires KAFKA_ADVERTISED_HOST \ + to be set to a reachable hostname or IP for Metadata broker advertisement" + .into(), + )); + } else { + local_addr.ip().to_string() + }; + + Ok(Self { host, port }) + } +} + +pub struct KafkaServer { + config: Arc<ServerConfig>, +} + +impl KafkaServer { + #[must_use] + pub fn new(config: ServerConfig) -> Self { + Self { + config: Arc::new(config), + } + } + + /// Accept Kafka wire connections until `shutdown` fires, then drain in-flight tasks. + /// + /// `listener` must already be bound by the caller. This lets tests and `main` bind + /// the port before spawning the task, eliminating the TOCTOU race of bind-drop-rebind. + /// + /// # Errors + /// + /// Returns an error on invalid config or a non-transient `accept()` error. + pub async fn run( + self, + listener: TcpListener, + mut shutdown: broadcast::Receiver<()>, + ) -> Result<()> { + let local_addr = listener.local_addr()?; + let broker = Arc::new(BrokerAdvertise::from_server_config( + &self.config, + local_addr, + )?); + info!( + "kafka listener bound on {} (advertised as {}:{})", + local_addr, broker.host, broker.port + ); + + let tracker = TaskTracker::new(); + let broker = Arc::clone(&broker); + + loop { + tokio::select! { + result = shutdown.recv() => { + match result { + Ok(()) => { + info!("kafka listener shutdown requested"); + tracker.close(); + tracker.wait().await; + break; + } + // Capacity-1 channel: lagged means a signal was sent before we polled — treat as shutdown. + Err(broadcast::error::RecvError::Lagged(_)) => { + info!("kafka listener shutdown requested (lagged)"); + tracker.close(); + tracker.wait().await; + break; + } + Err(broadcast::error::RecvError::Closed) => { + tracker.close(); + tracker.wait().await; + break; + } + } + } + accept_result = listener.accept() => { + match accept_result { + Ok((stream, peer)) => { + if let Err(e) = stream.set_nodelay(true) { + warn!(%peer, "TCP_NODELAY failed: {e}"); + } + let cfg = Arc::clone(&self.config); + let broker = Arc::clone(&broker); + tracker.spawn(async move { + if let Err(err) = handle_connection(stream, cfg, peer, broker).await { + warn!(%peer, "connection closed with error: {err}"); + } + }); + } + Err(e) if is_transient_accept_error(&e) => { + // Brief backoff on fd exhaustion to avoid busy-spinning. + if matches!(e.raw_os_error(), Some(23 | 24)) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + warn!(%e, "transient accept error, continuing"); + } + Err(e) => return Err(e.into()), + } + } + } + } + Ok(()) + } +} + +fn is_transient_accept_error(err: &std::io::Error) -> bool { + use std::io::ErrorKind; + + matches!( + err.kind(), + ErrorKind::Interrupted | ErrorKind::ConnectionAborted | ErrorKind::WouldBlock + ) || matches!( + err.raw_os_error(), + // EMFILE / ENFILE are common across Unix platforms when fd limits are hit. + Some(23 | 24) + ) +} + +async fn handle_connection( + mut stream: TcpStream, + config: Arc<ServerConfig>, + peer: SocketAddr, + broker: Arc<BrokerAdvertise>, +) -> Result<()> { + debug!(%peer, "connection accepted"); + + loop { + let frame = match read_frame(&mut stream, config.max_frame_size, config.read_timeout).await + { + Ok(f) => f, + Err(KafkaProtocolError::Io(ref e)) + if e.kind() == std::io::ErrorKind::UnexpectedEof + || e.kind() == std::io::ErrorKind::ConnectionReset => + { + info!(%peer, "connection closed by client"); + return Ok(()); + } + Err(e) => return Err(e), + }; + + if frame.len() < 8 { + return Err(KafkaProtocolError::BufferUnderflow { + needed: 8, + remaining: frame.len(), + }); + } + let api_key = i16::from_be_bytes([frame[0], frame[1]]); + let api_version = i16::from_be_bytes([frame[2], frame[3]]); + let req_hdr_ver = request_header_version(api_key, api_version); + let resp_hdr_ver = response_header_version(api_key, api_version); + let correlation_id = correlation_id_from_frame(&frame); + + let mut decoder = Decoder::new(frame); + let req = match RequestHeader::decode_from(&mut decoder, req_hdr_ver) { + Ok(req) => req, + Err(KafkaProtocolError::UnsupportedHeaderVersion(_)) => { Review Comment: Fixed in [9afa082](https://github.com/apache/iggy/pull/3519/commits/9afa082fcb6bba31ed0b754c0a71007c9066b9c3) -- 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]
