This is an automated email from the ASF dual-hosted git repository. krishvishal pushed a commit to branch kafka-gateway-perf in repository https://gitbox.apache.org/repos/asf/iggy.git
commit 91356b27333e88c8a6282cffdf12bde430a1e423 Author: Krishna Vishal <[email protected]> AuthorDate: Mon Oct 12 23:03:42 2026 +0530 perf(gateways): give each Kafka producer its own Iggy client Also set TCP no-delay on Iggy clients and cap Produce by decoded bytes. --- gateways/kafka/README.md | 20 +- gateways/kafka/docs/IDEMPOTENCE.md | 16 +- gateways/kafka/docs/SCOPE.md | 3 +- .../kafka/src/bridge/iggy_bridge/group_offsets.rs | 4 +- gateways/kafka/src/bridge/iggy_bridge/mod.rs | 23 +- gateways/kafka/src/bridge/iggy_bridge/produce.rs | 216 ++++++++++++++++- gateways/kafka/src/bridge/mod.rs | 4 +- gateways/kafka/src/protocol/api.rs | 51 +++- gateways/kafka/src/protocol/handlers/mod.rs | 2 +- gateways/kafka/src/protocol/handlers/produce.rs | 268 ++++++++++++++++++--- gateways/kafka/src/records.rs | 30 +++ gateways/kafka/tests/common/iggy_server.rs | 16 +- gateways/kafka/tests/produce_real_bridge_tests.rs | 3 +- 13 files changed, 580 insertions(+), 76 deletions(-) diff --git a/gateways/kafka/README.md b/gateways/kafka/README.md index 3b88850b7..9fbc1ad5a 100644 --- a/gateways/kafka/README.md +++ b/gateways/kafka/README.md @@ -356,7 +356,12 @@ Partition cap: `max_frame_size` decompressed bytes, `max_frame_size / 64` record per slot. Past it: 10. Request budget: 8 partition caps, refused partitions included. Past it: 6 for the rest, not -decoded. 4 requests decode at once. +decoded. + +Decode memory: 4 partition caps across all requests, counted in bytes. A request holds what its +largest partition may decode: its bytes or 64 per declared record, whichever is more, or a whole +cap when compressed. One holding 64 KiB or more first takes one of 4 large slots, and holds no +memory while it waits for one. Past its deadline waiting for either: 7. [docs/BRIDGE_MAPPING.md](docs/BRIDGE_MAPPING.md) describes what a record becomes once it is stored. @@ -507,13 +512,16 @@ cut before that threshold is reached - worth knowing rather than discovering lat ### Concurrency ceiling -- One `IggyClient` serves Produce, Metadata, CreateTopics, DescribeConfigs and AlterConfigs for - every Kafka connection, one Iggy request at a time. Fetch polls use 4 more, one per read slot. - Topic probes use 1 more. Offset calls use 8 more: a commit client and a read client per offset slot. +- Each producer sends on an `IggyClient` of its own, one send at a time, so its batches land in + order. A producer is its producer id, or its connection when its batches carry none. Up to 256 + of these stay open. Opening one closes the free ones idle for 10 minutes, then the least + recently used free one if 256 are open. A Produce request that finds all 256 busy answers 7. +- Metadata, CreateTopics, DeleteTopics, DescribeConfigs and AlterConfigs share 1 more client, one + Iggy request at a time. Fetch polls use 4 more, one per read slot. Topic probes use 1 more. + Offset calls use 8 more: a commit client and a read client per offset slot. - OffsetCommit and OffsetFetch cost one Iggy call per partition, spread over the offset slots. In a cluster, an offset commit client can also open one connection per node for stores. -- `IGGY_KAFKA_MAX_CONNECTIONS` does not change that. A Produce client pool is a TODO in - [docs/SCOPE.md](docs/SCOPE.md). +- `IGGY_KAFKA_MAX_CONNECTIONS` does not change that. - Order: set `max.in.flight.requests.per.connection=1`, or `retries=0`. Otherwise a retried batch can land after later ones. diff --git a/gateways/kafka/docs/IDEMPOTENCE.md b/gateways/kafka/docs/IDEMPOTENCE.md index 11e9d57b8..ac5ecb9ac 100644 --- a/gateways/kafka/docs/IDEMPOTENCE.md +++ b/gateways/kafka/docs/IDEMPOTENCE.md @@ -51,12 +51,9 @@ Kafka producer id therefore needs one bound session per producer. The request id is different. The transmute copies the header and replaces only `group`, `client`, `session` and `user_id`, so a caller-supplied request id survives to the `ClientTable`. -One session per producer is a connection pool keyed by producer id. That is the pool the -README's "Concurrency ceiling" section already owes before -[#3535](https://github.com/apache/iggy/issues/3535) and -[#3536](https://github.com/apache/iggy/issues/3536). - -Two SDK seams are missing for it: +One session per producer is a connection pool keyed by producer id. The gateway has that pool, for +throughput: see the README's "Concurrency ceiling". Retries are still not deduplicated, because two +SDK seams are missing: - a caller-chosen client id at build time. `ConsensusSession::with_client_id` is public, but both construction sites build with `ConsensusSession::new()` @@ -130,7 +127,7 @@ a birthday collision, not a remote one. The id is a pool key, not a dedup identity. Under the design above, the dedup identity is the session's own random client id, minted at register. The producer id only decides which connection serves a producer. Kafka still requires it to be unique across the cluster, which is what the -instance number buys. It also has to stay unique across a restart once the pool lands or +instance number buys. It also has to stay unique across a restart once dedup lands or Produce persists, because `producer_epoch` is always 0, so a replayed id is a replayed `(producer_id, producer_epoch)` pair a live producer may still hold. The counter therefore starts at the wall clock in milliseconds rather than at 0. A restarted gateway starts above every id its @@ -186,7 +183,7 @@ Produce: gateway closes the connection instead, as a Kafka broker does on any `acks=0` produce error. The check runs before the `acks=0` branch, so no write path can see a transactional batch -Those first two codes stay unsent even once the pool lands. The watermark accepts any request +Those first two codes stay unsent even once dedup lands. The watermark accepts any request above it without noticing a gap, so a gap cannot be told apart from ordinary traffic. Sending either code claims a detection the gateway does not have. @@ -235,8 +232,7 @@ A producer that holds an id believes its retries are deduplicated. They are not. network timeout writes the record twice, and both copies reach the stream with their own offsets. Iggy's own deduplication does not help, because it guards the other hop. Delivery through the -gateway is at-least-once until the pool lands, and at-least-once across a gateway restart after -that. +gateway is at-least-once until dedup lands, and at-least-once across a gateway restart after that. The README states that limitation under "Delivery guarantees", so a user does not have to infer it from the presence of key 22. diff --git a/gateways/kafka/docs/SCOPE.md b/gateways/kafka/docs/SCOPE.md index 91d2e7e36..71d813b71 100644 --- a/gateways/kafka/docs/SCOPE.md +++ b/gateways/kafka/docs/SCOPE.md @@ -193,7 +193,8 @@ below it are still open for the issues that build on top of it. - [x] Produce → `send_messages` ([#3535](https://github.com/apache/iggy/issues/3535)) - one call per partition, base offset from the send confirmation, one error code per partition. See [README.md](../README.md#produce-3535). -- [ ] Produce: `IggyClient` pool. Pin each partition to one client, so order holds. +- [x] Produce: `IggyClient` pool, one client per producer, so a producer's batches land in order. + See [README.md](../README.md#concurrency-ceiling) and [IDEMPOTENCE.md](IDEMPOTENCE.md). - [ ] Produce: write keyed records' header TLVs into one buffer (`records::to_iggy`). Benchmark first. Keep Iggy's TLV layout and the 100 KB header check. - [ ] Produce: skip the second batch CRC pass (`records::preflight`) only if profiling shows it diff --git a/gateways/kafka/src/bridge/iggy_bridge/group_offsets.rs b/gateways/kafka/src/bridge/iggy_bridge/group_offsets.rs index c455bce1e..ccb1677cd 100644 --- a/gateways/kafka/src/bridge/iggy_bridge/group_offsets.rs +++ b/gateways/kafka/src/bridge/iggy_bridge/group_offsets.rs @@ -412,7 +412,7 @@ mod tests { use tokio::sync::oneshot; use super::*; - use crate::bridge::iggy_bridge::FetchPool; + use crate::bridge::iggy_bridge::{FetchPool, ProducePool}; use crate::bridge::{DEFAULT_MAX_MESSAGE_SIZE, IggyBridgeConfig, TopicMapping}; fn target(topic: &str) -> TopicTarget { @@ -458,7 +458,7 @@ mod tests { topic_mapping: TopicMapping::new("kafka".to_string(), HashMap::new()).unwrap(), max_message_size: DEFAULT_MAX_MESSAGE_SIZE, }, - send_slot: Arc::new(Semaphore::new(1)), + produce_pool: ProducePool::new(), fetch_pool: Arc::new(FetchPool::new()), probe_client: LazyClient::default(), offset_pool: OffsetPool { diff --git a/gateways/kafka/src/bridge/iggy_bridge/mod.rs b/gateways/kafka/src/bridge/iggy_bridge/mod.rs index b0c8032ef..5ed0a7a4d 100644 --- a/gateways/kafka/src/bridge/iggy_bridge/mod.rs +++ b/gateways/kafka/src/bridge/iggy_bridge/mod.rs @@ -40,6 +40,8 @@ use fetch::{FetchPool, LazyClient}; pub(crate) use fetch::{FetchSlot, PartitionProbe, TopicProbe}; use group_offsets::OffsetPool; pub use group_offsets::{OFFSET_GROUP_PREFIX, OffsetCalls}; +use produce::ProducePool; +pub use produce::{ProduceSession, Producer}; pub use topics::{KafkaTopicMetadata, StreamTopicCache, TopicCreationOutcome, TopicLoad}; /// Passes attempted, after the first, before [`IggyBridge::connect`] gives up and returns `Err`. @@ -114,17 +116,19 @@ async fn permit_by(turn: &Arc<Semaphore>, deadline: Instant) -> Option<OwnedSema /// Owns the connected `IggyClient`s and resolves Kafka topics against them. /// -/// One lockstep client serves Produce, Metadata and `CreateTopics` for every Kafka connection, so -/// those Iggy calls run one at a time. Topic probes use a client of their own, Fetch polls one per -/// read slot, and offset calls two per offset slot. A Produce pool is a TODO in `docs/SCOPE.md`. +/// Each producer sends on a client of its own, one send at a time. Metadata, `CreateTopics`, +/// `DeleteTopics`, `DescribeConfigs` and `AlterConfigs` share one more client. Topic probes use a +/// client of their own, Fetch polls one per read slot, and offset calls two per offset slot. /// /// Each client signs in at the metadata leader, and a view change or a refused call can move it to /// another node. So in a cluster, two clients can read replicas that are not at the same offset. pub struct IggyBridge { + /// The client of Metadata, `CreateTopics`, `DeleteTopics`, `DescribeConfigs` and + /// `AlterConfigs`. client: Arc<IggyClient>, config: IggyBridgeConfig, - /// One Produce send inside the SDK at a time. See `send_records`. - send_slot: Arc<Semaphore>, + /// One Iggy session per producer. See `send_records`. + produce_pool: ProducePool, /// Fetch read slots, each with its own client. See `fetch_slot`. fetch_pool: Arc<FetchPool>, /// The client of every topic probe, Fetch's and `ListOffsets`'. See `probe`. @@ -165,7 +169,7 @@ impl IggyBridge { Ok(Self { client: Arc::new(client), config, - send_slot: Arc::new(Semaphore::new(1)), + produce_pool: ProducePool::new(), fetch_pool: Arc::new(FetchPool::new()), probe_client: LazyClient::default(), offset_pool: OffsetPool::new(), @@ -222,11 +226,13 @@ impl IggyBridge { /// was already in a state that rejects a clean shutdown). Every client is shut down even when /// one fails, and the error is the first one. pub async fn close(self) -> Result<(), BridgeError> { + let produce_clients = self.produce_pool.close().await; let fetch_clients = self.fetch_pool.close().await; let probe_client = self.probe_client.close().await; let offset_clients = self.offset_pool.close().await; let shared_client = with_request_timeout(self.client.shutdown()).await; - fetch_clients + produce_clients + .and(fetch_clients) .and(probe_client) .and(offset_clients) .and(shared_client) @@ -243,9 +249,12 @@ async fn connect_client(config: &IggyBridgeConfig) -> Result<IggyClient, BridgeE let credentials = Credentials::UsernamePassword(config.username.clone(), config.password.clone()); + // A send over the SDK's 8 KiB write buffer leaves as two writes, header then payload. With + // Nagle on, the payload waits for the server's delayed ACK of the header: about 40 ms a send. let client = IggyClientBuilder::new() .with_tcp() .with_server_address(config.address.clone()) + .with_no_delay() .with_auto_sign_in(AutoLogin::Enabled(credentials)) .with_reconnection_max_retries(Some(RECONNECTION_RETRIES)) .build() diff --git a/gateways/kafka/src/bridge/iggy_bridge/produce.rs b/gateways/kafka/src/bridge/iggy_bridge/produce.rs index b4ee4d09e..dba3c1c41 100644 --- a/gateways/kafka/src/bridge/iggy_bridge/produce.rs +++ b/gateways/kafka/src/bridge/iggy_bridge/produce.rs @@ -19,19 +19,122 @@ //! //! Separate from `fetch.rs` so the two never edit one file. +use std::collections::HashMap; use std::future::Future; -use std::sync::Arc; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::Duration; use iggy::prelude::{IggyError, IggyMessage, MessageClient, Partitioning}; use tokio::sync::Semaphore; -use tokio::time::{Instant, timeout}; +use tokio::time::{Instant, timeout, timeout_at}; +use super::fetch::LazyClient; use super::{IggyBridge, SLOT_LIMIT, TopicTarget, in_slot, permit_by}; use crate::bridge::error::BridgeError; /// Frame bytes besides the messages: request header, batch header and ids. const SEND_FRAME_ROOM: u64 = 4 * 1024; +/// Produce sessions open at once, each an Iggy connection. +const PRODUCE_SESSIONS: usize = 256; + +/// A free session unused this long closes when a new producer opens one. +const SESSION_IDLE: Duration = Duration::from_secs(600); + +/// Who sends a Produce request. Each producer gets an Iggy session of its own. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Producer { + /// The producer id of the request's record batches. + Id(i64), + /// The Kafka connection, for a producer without an id. + Connection(u64), +} + +/// One producer's Iggy client, and one send at a time on it, so its sends reach Iggy in order. +pub struct ProduceSession { + client: LazyClient, + turn: Arc<Semaphore>, +} + +/// The Produce sessions, one per producer, so different producers send at once. +pub(super) struct ProducePool { + sessions: Mutex<HashMap<Producer, Held>>, + capacity: usize, +} + +/// A session, and when a request last took it. +struct Held { + session: Arc<ProduceSession>, + used: Instant, +} + +impl ProducePool { + pub(super) fn new() -> Self { + Self::with_capacity(PRODUCE_SESSIONS) + } + + fn with_capacity(capacity: usize) -> Self { + Self { + sessions: Mutex::new(HashMap::new()), + capacity, + } + } + + /// The session of `producer`, opened if it has none. `None` when the pool is full and every + /// session is busy. + /// + /// Opening one first drops the free sessions idle past [`SESSION_IDLE`], then the least + /// recently used free one if the pool is still full. + fn take(&self, producer: Producer) -> Option<Arc<ProduceSession>> { + let now = Instant::now(); + let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(held) = sessions.get_mut(&producer) { + held.used = now; + return Some(Arc::clone(&held.session)); + } + sessions + .retain(|_, held| !(held.is_free() && now.duration_since(held.used) >= SESSION_IDLE)); + if sessions.len() >= self.capacity { + let oldest = sessions + .iter() + .filter(|(_, held)| held.is_free()) + .min_by_key(|(_, held)| held.used) + .map(|(producer, _)| *producer)?; + sessions.remove(&oldest); + } + let session = Arc::new(ProduceSession { + client: LazyClient::default(), + turn: Arc::new(Semaphore::new(1)), + }); + let held = Held { + session: Arc::clone(&session), + used: now, + }; + sessions.insert(producer, held); + drop(sessions); + Some(session) + } + + /// Shuts down each session's client that connected. See [`LazyClient::close_all`]. + pub(super) async fn close(&self) -> Result<(), BridgeError> { + let clients: Vec<LazyClient> = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .values() + .map(|held| held.session.client.clone()) + .collect(); + LazyClient::close_all(&clients).await + } +} + +impl Held { + /// No request holds it and no send runs on it, so dropping it loses nothing. + fn is_free(&self) -> bool { + Arc::strong_count(&self.session) == 1 && self.session.turn.available_permits() == 1 + } +} + impl IggyBridge { /// Largest partition, in message bytes, that one Iggy send carries. #[must_use] @@ -39,7 +142,13 @@ impl IggyBridge { self.config.max_message_size.saturating_sub(SEND_FRAME_ROOM) } - /// Appends `messages` to `partition` of `target` and returns the base offset. + /// The Produce session of `producer`. `None` when every session is busy. + #[must_use] + pub fn produce_session(&self, producer: Producer) -> Option<Arc<ProduceSession>> { + self.produce_pool.take(producer) + } + + /// Appends `messages` to `partition` of `target` on `session`, and returns the base offset. /// /// `partition` passes through unchanged: both systems count from 0. The server rejects a /// partition the topic lacks, so no round trip checks first. @@ -47,22 +156,26 @@ impl IggyBridge { /// `Ok(None)`: committed, but the server named no offset (a duplicate request, or a journal /// entry that is gone). Do not retry it. /// - /// A stuck send holds the slot up to `SLOT_LIMIT`. Other sends answer 7 meanwhile. + /// A stuck send holds its session up to `SLOT_LIMIT`, so the producer's later sends answer 7 + /// meanwhile and cannot overtake it. Other producers send on. /// /// # Errors /// /// [`BridgeError::Timeout`] if `deadline` passes. The send may have landed. /// [`BridgeError::SendLost`] if the connection broke during the send. It may have landed. - /// [`BridgeError::Iggy`] if the stream, topic or partition is missing, or the send is too - /// large. + /// [`BridgeError::Iggy`] if the session cannot connect, if the stream, topic or partition is + /// missing, or if the send is too large. pub async fn send_records( &self, + session: &ProduceSession, target: &TopicTarget, partition: u32, mut messages: Vec<IggyMessage>, deadline: Instant, ) -> Result<Option<u64>, BridgeError> { - let client = Arc::clone(&self.client); + let client = timeout_at(deadline, session.client.connected(&self.config)) + .await + .map_err(|_elapsed| BridgeError::Timeout)??; let (stream_id, topic_id) = (target.stream_id.clone(), target.topic_id.clone()); let send = async move { let partitioning = Partitioning::partition_id(partition); @@ -70,7 +183,7 @@ impl IggyBridge { .send_messages(&stream_id, &topic_id, &partitioning, &mut messages) .await }; - let response = send_in_slot(&self.send_slot, deadline, send).await?; + let response = send_in_slot(&session.turn, deadline, send).await?; Ok(response .confirmations .iter() @@ -191,4 +304,91 @@ mod tests { Err(BridgeError::SendLost(IggyError::Disconnected)) )); } + + #[tokio::test] + async fn given_one_producer_when_taken_twice_should_get_one_session() { + let pool = ProducePool::with_capacity(4); + let first = pool.take(Producer::Id(7)).unwrap(); + + assert!(Arc::ptr_eq(&first, &pool.take(Producer::Id(7)).unwrap())); + assert!( + !Arc::ptr_eq(&first, &pool.take(Producer::Connection(7)).unwrap()), + "an id and a connection number never share a session" + ); + } + + #[tokio::test] + async fn given_a_stuck_send_when_another_producer_sends_should_not_wait_for_it() { + let pool = ProducePool::with_capacity(4); + let (stuck, other) = ( + pool.take(Producer::Id(1)).unwrap(), + pool.take(Producer::Id(2)).unwrap(), + ); + let (_release, released) = oneshot::channel::<()>(); + let stuck_send = async move { + let _ = released.await; + Ok::<u64, IggyError>(1) + }; + assert!(matches!( + send_in_slot(&stuck.turn, Instant::now() + WAIT, stuck_send).await, + Err(BridgeError::Timeout) + )); + + let started = Arc::new(AtomicBool::new(false)); + let sent = send_in_slot(&other.turn, Instant::now() + WAIT * 10, flagged(started)).await; + assert_eq!(sent.unwrap(), 2); + } + + #[tokio::test(start_paused = true)] + async fn given_a_full_pool_when_a_new_producer_comes_should_drop_the_least_used_free_one() { + let pool = ProducePool::with_capacity(2); + let first = Arc::downgrade(&pool.take(Producer::Id(1)).unwrap()); + tokio::time::advance(WAIT).await; + let second = Arc::downgrade(&pool.take(Producer::Id(2)).unwrap()); + tokio::time::advance(WAIT).await; + pool.take(Producer::Id(1)).unwrap(); + tokio::time::advance(WAIT).await; + + pool.take(Producer::Id(3)).unwrap(); + assert!(first.upgrade().is_some(), "used most recently, so it stays"); + assert!(second.upgrade().is_none()); + } + + #[tokio::test] + async fn given_a_full_pool_of_held_sessions_when_a_new_producer_comes_should_refuse_it() { + let pool = ProducePool::with_capacity(1); + let held = pool.take(Producer::Id(1)).unwrap(); + + assert!(pool.take(Producer::Id(2)).is_none()); + drop(held); + assert!( + pool.take(Producer::Id(2)).is_some(), + "a free session makes room" + ); + } + + #[tokio::test] + async fn given_a_send_in_flight_when_the_pool_is_full_should_keep_its_session() { + let pool = ProducePool::with_capacity(1); + let session = pool.take(Producer::Id(1)).unwrap(); + let _sending = Arc::clone(&session.turn).acquire_owned().await.unwrap(); + let sending = Arc::downgrade(&session); + drop(session); + + assert!( + pool.take(Producer::Id(2)).is_none(), + "the producer's next send queues behind this one, not on a new session" + ); + assert!(sending.upgrade().is_some()); + } + + #[tokio::test(start_paused = true)] + async fn given_a_session_idle_past_the_limit_when_a_new_producer_comes_should_drop_it() { + let pool = ProducePool::with_capacity(4); + let idle = Arc::downgrade(&pool.take(Producer::Connection(1)).unwrap()); + tokio::time::advance(SESSION_IDLE).await; + + pool.take(Producer::Id(2)).unwrap(); + assert!(idle.upgrade().is_none()); + } } diff --git a/gateways/kafka/src/bridge/mod.rs b/gateways/kafka/src/bridge/mod.rs index 6083448f4..cfa34a60d 100644 --- a/gateways/kafka/src/bridge/mod.rs +++ b/gateways/kafka/src/bridge/mod.rs @@ -29,7 +29,7 @@ pub mod topic_map; pub use config::{DEFAULT_MAX_MESSAGE_SIZE, IggyBridgeConfig}; pub use error::BridgeError; pub use iggy_bridge::{ - IggyBridge, KafkaTopicMetadata, OFFSET_GROUP_PREFIX, OffsetCalls, StreamTopicCache, - TopicCreationOutcome, TopicLoad, TopicTarget, + IggyBridge, KafkaTopicMetadata, OFFSET_GROUP_PREFIX, OffsetCalls, ProduceSession, Producer, + StreamTopicCache, TopicCreationOutcome, TopicLoad, TopicTarget, }; pub use topic_map::{TopicMapping, TopicOverride}; diff --git a/gateways/kafka/src/protocol/api.rs b/gateways/kafka/src/protocol/api.rs index 9539773b8..33a878351 100644 --- a/gateways/kafka/src/protocol/api.rs +++ b/gateways/kafka/src/protocol/api.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; @@ -318,9 +319,9 @@ pub fn supported_api_ranges() -> &'static [ApiVersionRange] { /// `bridge` is `None` until `IGGY_KAFKA_BRIDGE_ENABLED` turns it on. A handler that finds `None` /// answers with its stub, so APIs can be wired one at a time. /// -/// Every Iggy call but a Fetch poll, a topic probe and an offset call goes through one lockstep -/// `IggyClient`, so Kafka connections serialize behind whichever of those calls is in flight. The -/// `Arc` does not change that. See the README's "Concurrency ceiling". +/// Produce sends on one Iggy client per producer. Metadata, `CreateTopics`, `DeleteTopics`, +/// `DescribeConfigs` and `AlterConfigs` share one lockstep `IggyClient`, so those calls run one +/// at a time across connections. See the README's "Concurrency ceiling". pub struct GatewayState { pub broker: BrokerAdvertise, pub bridge: Option<Arc<IggyBridge>>, @@ -331,8 +332,11 @@ pub struct GatewayState { /// Shared across every connection this gateway serves: a producer id has to be unique for /// the process, not for the connection that asked for it. pub producer_ids: ProducerIdAllocator, - /// Produce requests that decode and send at once. Caps their memory. - pub(crate) produce_slots: Semaphore, + /// `PRODUCE_CAPS` partition caps, in frame bytes, that Produce may hold decoded at once. Caps + /// its memory. A request takes what its largest partition may decode. + pub(crate) produce_memory: Semaphore, + /// `PRODUCE_LARGE` slots a large Produce request takes before its memory. Small ones skip it. + pub(crate) produce_large: Semaphore, /// Consumer group membership. Process-wide and independent of the bridge: a member outlives /// the connection that created it, and group coordination needs no Iggy call. pub groups: GroupCoordinator, @@ -343,9 +347,15 @@ pub struct GatewayState { pub(crate) probe_board: ProbeBoard, } -/// Each holds one decoded partition at a time, so about 160 MB at the default 8 MiB frame. Sends -/// run one at a time anyway. -const PRODUCE_SLOTS: usize = 4; +/// Partition caps Produce may hold decoded at once, about 160 MB at the default 8 MiB frame. +/// +/// Counted in bytes, not requests: a small request takes a small share, so it does not queue +/// behind a few large ones. +const PRODUCE_CAPS: usize = 4; + +/// Large Produce requests that decode and send at once. More at once only contend for the +/// gateway's workers and for Iggy: at full load they cost CPU and gave no throughput. +const PRODUCE_LARGE: usize = 4; impl GatewayState { #[must_use] @@ -363,7 +373,12 @@ impl GatewayState { max_frame_size, sasl_enabled, producer_ids: ProducerIdAllocator::new(instance_id), - produce_slots: Semaphore::const_new(PRODUCE_SLOTS), + produce_memory: Semaphore::new( + PRODUCE_CAPS + .saturating_mul(max_frame_size) + .min(Semaphore::MAX_PERMITS), + ), + produce_large: Semaphore::const_new(PRODUCE_LARGE), groups, loading: Mutex::default(), probe_board: ProbeBoard::default(), @@ -400,9 +415,27 @@ pub struct ConnectionState { /// The member that last sent a Heartbeat here. An offset call caps its hold by that member's /// session, since its next heartbeat waits behind the call. heartbeat_member: Mutex<Option<GroupMember>>, + /// Keys the Produce session of a producer that sends no producer id. + id: ConnectionId, +} + +/// A Kafka connection's number, unique in the process. +#[derive(Debug, Clone, Copy)] +struct ConnectionId(u64); + +static NEXT_CONNECTION_ID: AtomicU64 = AtomicU64::new(0); + +impl Default for ConnectionId { + fn default() -> Self { + Self(NEXT_CONNECTION_ID.fetch_add(1, Ordering::Relaxed)) + } } impl ConnectionState { + pub(crate) const fn id(&self) -> u64 { + self.id.0 + } + pub(crate) fn note_heartbeat(&self, member: GroupMember) { *self .heartbeat_member diff --git a/gateways/kafka/src/protocol/handlers/mod.rs b/gateways/kafka/src/protocol/handlers/mod.rs index 2e2282229..2b86a5f7e 100644 --- a/gateways/kafka/src/protocol/handlers/mod.rs +++ b/gateways/kafka/src/protocol/handlers/mod.rs @@ -111,7 +111,7 @@ pub async fn dispatch( body: Bytes, ) -> HandleOutcome { match api_key { - API_KEY_PRODUCE => produce::handle(state, api_version, body).await, + API_KEY_PRODUCE => produce::handle(state, connection, api_version, body).await, API_KEY_FETCH => fetch::handle(state, connection, api_version, body).await, API_KEY_LIST_OFFSETS => list_offsets::handle(state, api_version, body).await, API_KEY_METADATA => metadata::handle(state, api_version, body).await, diff --git a/gateways/kafka/src/protocol/handlers/produce.rs b/gateways/kafka/src/protocol/handlers/produce.rs index a18b9c388..1a1421014 100644 --- a/gateways/kafka/src/protocol/handlers/produce.rs +++ b/gateways/kafka/src/protocol/handlers/produce.rs @@ -25,13 +25,14 @@ use kafka_protocol::messages::produce_request::PartitionProduceData; use kafka_protocol::messages::produce_response::{PartitionProduceResponse, TopicProduceResponse}; use kafka_protocol::messages::{ProduceRequest, ProduceResponse}; use kafka_protocol::protocol::StrBytes; +use tokio::sync::{AcquireError, SemaphorePermit}; -use crate::bridge::{BridgeError, IggyBridge, TopicTarget}; +use crate::bridge::{BridgeError, IggyBridge, ProduceSession, Producer, TopicTarget}; use crate::error::Result; use crate::protocol::api::{ - API_KEY_PRODUCE, ApiVersionRange, ERROR_INVALID_RECORD, ERROR_INVALID_REQUIRED_ACKS, - ERROR_MESSAGE_TOO_LARGE, ERROR_NONE, ERROR_NOT_LEADER_OR_FOLLOWER, ERROR_REQUEST_TIMED_OUT, - ERROR_UNKNOWN_SERVER_ERROR, ERROR_UNKNOWN_TOPIC_OR_PARTITION, + API_KEY_PRODUCE, ApiVersionRange, ConnectionState, ERROR_INVALID_RECORD, + ERROR_INVALID_REQUIRED_ACKS, ERROR_MESSAGE_TOO_LARGE, ERROR_NONE, ERROR_NOT_LEADER_OR_FOLLOWER, + ERROR_REQUEST_TIMED_OUT, ERROR_UNKNOWN_SERVER_ERROR, ERROR_UNKNOWN_TOPIC_OR_PARTITION, ERROR_UNSUPPORTED_COMPRESSION_TYPE, ERROR_UNSUPPORTED_VERSION, GatewayState, HandleOutcome, REQUEST_DEADLINE, UNKNOWN_OFFSET, supported_max_version, }; @@ -42,8 +43,8 @@ use crate::protocol::handlers::{ }; use crate::protocol::probe_board::ProbeBoard; use crate::records::{ - Allowance, DecompressionBudget, RecordCodecError, TimestampWindow, Zstd, decode_batch, - is_compressed, to_iggy, + Allowance, DecompressionBudget, RecordCodecError, TimestampWindow, Zstd, declared_records, + decode_batch, is_compressed, producer_id, to_iggy, }; pub const RANGE: ApiVersionRange = ApiVersionRange { @@ -63,7 +64,7 @@ const ZSTD_MIN_VERSION: i16 = 7; /// Frame bytes per record slot. Caps one partition at 131,072 slots, about 40 MB, at the default /// 8 MiB frame. /// -/// `GatewayState` caps how many requests hold that at once. +/// `GatewayState` caps how much of that requests hold at once, at the same rate. const FRAME_BYTES_PER_RECORD: usize = 64; /// Partition allowances one request may decode, refused partitions included. @@ -82,7 +83,12 @@ const RESPONSE_BASE_BYTES: usize = 512; /// spec-compliant client can legitimately send Produce v0-2 with `acks=0`. Rejecting those /// versions before reading `acks` would send an error response the client never expects, /// desyncing the next correlation id it reads. -pub async fn handle(state: &GatewayState, api_version: i16, body: Bytes) -> HandleOutcome { +pub async fn handle( + state: &GatewayState, + connection: &ConnectionState, + api_version: i16, + body: Bytes, +) -> HandleOutcome { // Above the encoder max there is no response parseable at the client's version, so close // rather than decode (same policy the other APIs apply). Fail-closed on a missing row // (i16::MIN, not i16::MAX): dispatch routes Produce on its api key, not from this table, so @@ -156,17 +162,14 @@ pub async fn handle(state: &GatewayState, api_version: i16, body: Bytes) -> Hand // Before planning, not after: the client waits through decompression too. let deadline = Deadline::new(request.timeout_ms); - let Ok(Ok(_slot)) = tokio::time::timeout_at(deadline.at, state.produce_slots.acquire()).await - else { - tracing::debug!("Produce deadline passed while waiting for a slot"); - return if request.acks == ACKS_NONE { - HandleOutcome::Close - } else { - respond_or_close( - encode_uniform_response(api_version, &request, ERROR_REQUEST_TIMED_OUT), - "Produce", - ) - }; + let cost = decode_cost(&request, state.max_frame_size); + let Ok(Ok(_admission)) = tokio::time::timeout_at(deadline.at, admit(state, cost)).await else { + tracing::debug!("Produce deadline passed while waiting to decode"); + return timed_out(api_version, &request); + }; + let Some(session) = bridge.produce_session(producer_of(&request, connection)) else { + tracing::warn!("every Produce session is busy"); + return timed_out(api_version, &request); }; let partition = Allowance { bytes: state.max_frame_size, @@ -178,8 +181,8 @@ pub async fn handle(state: &GatewayState, api_version: i16, body: Bytes) -> Hand } else { Zstd::Refused }; - let responses = - write_request(bridge, &state.probe_board, &request, budget, zstd, deadline).await; + let board = &state.probe_board; + let responses = write_request(bridge, &session, board, &request, budget, zstd, deadline).await; // Kafka closes an `acks=0` connection on any error, so the client refreshes metadata. if request.acks == ACKS_NONE { @@ -206,6 +209,71 @@ fn stub_outcome(api_version: i16, request: &ProduceRequest) -> HandleOutcome { respond_or_close(encode_response(api_version, request), "Produce") } +/// Answers 7 for every partition. `acks=0` has no response to carry it, so that closes instead. +fn timed_out(api_version: i16, request: &ProduceRequest) -> HandleOutcome { + if request.acks == ACKS_NONE { + HandleOutcome::Close + } else { + respond_or_close( + encode_uniform_response(api_version, request, ERROR_REQUEST_TIMED_OUT), + "Produce", + ) + } +} + +/// What the request may hold decoded, in frame bytes, at most one partition cap. +/// +/// Partitions decode one at a time, so the largest counts. Uncompressed, a partition holds its own +/// bytes or a record slot per declared record, whichever is more. Compressed, it may fill the cap. +fn decode_cost(request: &ProduceRequest, max_frame_size: usize) -> usize { + let largest = request + .topic_data + .iter() + .flat_map(|topic| &topic.partition_data) + .filter_map(|partition| partition.records.as_deref()) + .map(|blob| { + if is_compressed(blob) { + max_frame_size + } else { + let slots = declared_records(blob).unwrap_or(0); + blob.len().max(slots.saturating_mul(FRAME_BYTES_PER_RECORD)) + } + }) + .max() + .unwrap_or(0); + largest.min(max_frame_size) +} + +/// A large slot when the request is large, and its decode memory. +type Admission<'a> = (Option<SemaphorePermit<'a>>, SemaphorePermit<'a>); + +/// Waits for the request's decode memory. A large request, `CODEC_OFF_WORKER_BYTES` or more, +/// first waits for a large slot, so while it waits its turn it holds no memory. +async fn admit( + state: &GatewayState, + cost: usize, +) -> std::result::Result<Admission<'_>, AcquireError> { + let large = if cost >= CODEC_OFF_WORKER_BYTES { + Some(state.produce_large.acquire().await?) + } else { + None + }; + let permits = u32::try_from(cost).unwrap_or(u32::MAX); + let memory = state.produce_memory.acquire_many(permits).await?; + Ok((large, memory)) +} + +/// The producer id of the request's first batch, or the connection when that batch has none. +fn producer_of(request: &ProduceRequest, connection: &ConnectionState) -> Producer { + request + .topic_data + .iter() + .flat_map(|topic| &topic.partition_data) + .find_map(|partition| partition.records.as_deref()) + .and_then(producer_id) + .map_or_else(|| Producer::Connection(connection.id()), Producer::Id) +} + /// When a request stops writing. #[derive(Clone, Copy)] struct Deadline { @@ -258,10 +326,10 @@ impl From<RecordCodecError> for Refusal { /// Only one partition's messages live at a time. The budget is owned here and borrowed only in /// sync calls: a borrow held across an await makes the connection task `!Send`. /// -/// Sequential: the one `IggyClient` behind the bridge is lockstep, so concurrent sends would -/// just queue on its stream mutex (see the README's "Concurrency ceiling"). +/// Sequential: the producer's session sends one at a time, so its sends reach Iggy in order. async fn write_request( bridge: &IggyBridge, + session: &ProduceSession, board: &ProbeBoard, request: &ProduceRequest, budget: DecompressionBudget, @@ -282,9 +350,11 @@ async fn write_request( Ok(target) => { let planned = plan(&budget, zstd, name, partition, max_send); match planned { - Ok((id, messages)) => send_by(bridge, target, name, id, messages, deadline) - .await - .map_err(Refusal::from), + Ok((id, messages)) => { + send_by(bridge, session, target, name, id, messages, deadline) + .await + .map_err(Refusal::from) + } Err(refusal) => Err(refusal), } } @@ -413,6 +483,7 @@ fn request_timeout(timeout_ms: i32) -> Duration { /// 7 means nothing was written, or the send may still land and a retry duplicate it. async fn send_by( bridge: &IggyBridge, + session: &ProduceSession, target: &TopicTarget, kafka_topic: &str, partition_id: u32, @@ -420,7 +491,7 @@ async fn send_by( deadline: Deadline, ) -> std::result::Result<Option<u64>, i16> { match bridge - .send_records(target, partition_id, messages, deadline.at) + .send_records(session, target, partition_id, messages, deadline.at) .await { Ok(base_offset) => Ok(base_offset), @@ -571,11 +642,14 @@ mod tests { use super::*; use bytes::BytesMut; use kafka_protocol::indexmap::IndexMap; + use kafka_protocol::messages::TopicName; + use kafka_protocol::messages::produce_request::TopicProduceData; use kafka_protocol::records::{ Compression, NO_PARTITION_LEADER_EPOCH, NO_PRODUCER_EPOCH, NO_PRODUCER_ID, NO_SEQUENCE, Record, RecordBatchEncoder, RecordEncodeOptions, TimestampType, }; + use crate::protocol::api::BrokerAdvertise; use crate::records::encode_batch; const CREATE_TIME: i64 = 1_700_000_000_123; @@ -694,6 +768,146 @@ mod tests { ); } + fn request_of(partitions: Vec<PartitionProduceData>) -> ProduceRequest { + ProduceRequest::default().with_topic_data(vec![ + TopicProduceData::default() + .with_name(TopicName(StrBytes::from_static_str(TOPIC))) + .with_partition_data(partitions), + ]) + } + + #[test] + fn given_an_idempotent_batch_when_keyed_should_take_its_producer_id() { + let mut records = [record(0, b"v")]; + records[0].producer_id = 7; + records[0].producer_epoch = 0; + records[0].sequence = 0; + let request = request_of(vec![entry( + 0, + Some(compressed_batch(&records, Compression::None)), + )]); + + let producer = producer_of(&request, &ConnectionState::default()); + assert_eq!(producer, Producer::Id(7)); + } + + #[test] + fn given_a_batch_without_a_producer_id_when_keyed_should_take_the_connection() { + let batch = compressed_batch(&[record(0, b"v")], Compression::None); + let request = request_of(vec![entry(0, Some(batch))]); + let connection = ConnectionState::default(); + + assert_eq!( + producer_of(&request, &connection), + Producer::Connection(connection.id()) + ); + assert_ne!( + ConnectionState::default().id(), + connection.id(), + "each connection keys a session of its own" + ); + } + + /// The default frame. + const FRAME: usize = 8 * 1024 * 1024; + /// Long enough for an admission that is free to complete. Time is paused in these tests. + const WAIT: Duration = Duration::from_millis(50); + + fn cost_of(partitions: Vec<PartitionProduceData>, max_frame_size: usize) -> usize { + decode_cost(&request_of(partitions), max_frame_size) + } + + #[test] + fn given_an_uncompressed_batch_when_costed_should_take_its_bytes() { + let batch = encode_batch(&mut [record(0, &[7; 1000]), record(1, &[7; 1000])]).unwrap(); + assert_eq!( + cost_of(vec![entry(0, Some(batch.clone()))], FRAME), + batch.len() + ); + } + + #[test] + fn given_small_records_when_costed_should_take_a_record_slot_each() { + let mut records: Vec<Record> = (0..100).map(|offset| record(offset, b"v")).collect(); + let batch = encode_batch(&mut records).unwrap(); + assert!(batch.len() < 100 * FRAME_BYTES_PER_RECORD); + assert_eq!( + cost_of(vec![entry(0, Some(batch))], FRAME), + 100 * FRAME_BYTES_PER_RECORD + ); + } + + #[test] + fn given_a_compressed_batch_when_costed_should_take_a_whole_cap() { + let batch = gzip_batch(&[record(0, b"v")]); + assert_eq!( + cost_of(vec![entry(0, Some(batch))], FRAME), + FRAME, + "it may inflate to the cap" + ); + } + + #[test] + fn given_several_partitions_when_costed_should_take_the_largest() { + let small = compressed_batch(&[record(0, &[7; 100])], Compression::None); + let large = compressed_batch(&[record(0, &[7; 5000])], Compression::None); + let partitions = vec![ + entry(0, Some(small)), + entry(1, Some(large.clone())), + entry(2, None), + ]; + assert_eq!( + cost_of(partitions, FRAME), + large.len(), + "partitions decode one at a time" + ); + } + + #[test] + fn given_a_partition_past_the_cap_when_costed_should_take_one_cap() { + let batch = compressed_batch(&[record(0, &[7; 5000])], Compression::None); + assert_eq!(cost_of(vec![entry(0, Some(batch))], 4096), 4096); + } + + /// Admits four large requests, which takes every large slot. + async fn large_slots_taken(state: &GatewayState) -> Vec<Admission<'_>> { + let mut admitted = Vec::new(); + for _ in 0..4 { + admitted.push(admit(state, CODEC_OFF_WORKER_BYTES).await.unwrap()); + } + admitted + } + + #[tokio::test(start_paused = true)] + async fn given_the_large_slots_taken_when_a_small_request_comes_should_admit_it() { + let state = GatewayState::stub(BrokerAdvertise::default(), FRAME); + let _large = large_slots_taken(&state).await; + + let fifth = tokio::time::timeout(WAIT, admit(&state, CODEC_OFF_WORKER_BYTES)).await; + assert!(fifth.is_err(), "a fifth large request waits its turn"); + let small = tokio::time::timeout(WAIT, admit(&state, CODEC_OFF_WORKER_BYTES - 1)).await; + assert!( + small.is_ok(), + "a small one does not queue behind large ones" + ); + } + + #[tokio::test(start_paused = true)] + async fn given_the_large_slots_taken_when_a_large_request_waits_should_hold_no_memory() { + let state = GatewayState::stub(BrokerAdvertise::default(), FRAME); + let _large = large_slots_taken(&state).await; + let free = state.produce_memory.available_permits(); + + let waiting = admit(&state, FRAME); + tokio::pin!(waiting); + assert!(tokio::time::timeout(WAIT, &mut waiting).await.is_err()); + assert_eq!( + state.produce_memory.available_permits(), + free, + "small requests can still take that memory" + ); + } + #[test] fn given_an_idempotent_batch_when_planned_should_convert_it() { let mut records = [record(0, b"v")]; diff --git a/gateways/kafka/src/records.rs b/gateways/kafka/src/records.rs index 42c8fbd51..db0572177 100644 --- a/gateways/kafka/src/records.rs +++ b/gateways/kafka/src/records.rs @@ -867,6 +867,20 @@ pub fn is_compressed(blob: &[u8]) -> bool { .is_some_and(|attributes| attributes & 0b111 != 0) } +/// The producer id of the first batch in `blob`, bytes 43 to 50. `None` when it has none. +#[must_use] +pub fn producer_id(blob: &[u8]) -> Option<i64> { + let id = i64::from_be_bytes(blob.get(43..51)?.try_into().ok()?); + (id != NO_PRODUCER_ID).then_some(id) +} + +/// The record count the first batch in `blob` declares, bytes 57 to 60. `None` when it has none. +#[must_use] +pub fn declared_records(blob: &[u8]) -> Option<usize> { + let count = i32::from_be_bytes(blob.get(57..BATCH_HEADER_BYTES)?.try_into().ok()?); + usize::try_from(count).ok() +} + /// Decodes the one record batch a Produce partition carries. /// /// # Errors @@ -1839,6 +1853,22 @@ mod tests { assert!(!is_compressed(&[])); } + #[test] + fn given_a_batch_when_its_records_are_counted_should_read_the_declared_count() { + let mut records = [record_at_offset(0, b"a"), record_at_offset(1, b"b")]; + assert_eq!( + declared_records(&encode_batch(&mut records).unwrap()), + Some(2) + ); + // `encode_batch` set the sequences that keep both records in one batch. + assert_eq!( + declared_records(&encode_with(&records, Compression::Gzip)), + Some(2), + "the header counts records before compression" + ); + assert_eq!(declared_records(&[0; BATCH_HEADER_BYTES - 1]), None); + } + #[test] fn given_a_gzip_batch_when_decoded_should_read_it() { let records = vec![record_at_offset(0, b"compressed")]; diff --git a/gateways/kafka/tests/common/iggy_server.rs b/gateways/kafka/tests/common/iggy_server.rs index 51cc9a746..82d320bc9 100644 --- a/gateways/kafka/tests/common/iggy_server.rs +++ b/gateways/kafka/tests/common/iggy_server.rs @@ -239,7 +239,8 @@ pub struct TestServer { child: Child, address: String, password: String, - _port_guard: PortGuard, + /// `None` once [`Self::stop`] handed it to the caller. + port_guard: Option<PortGuard>, } impl TestServer { @@ -306,12 +307,23 @@ impl TestServer { child, address, password: password.to_string(), - _port_guard: port_guard, + port_guard: Some(port_guard), }; server.wait_ready().await; server } + /// Stops the server and keeps its port until the returned guard drops. Without it, another + /// test's server can take the port while a client of this one still retries there. + pub fn stop(mut self) -> PortGuard { + let port_guard = self + .port_guard + .take() + .expect("a running server holds its port"); + drop(self); + port_guard + } + /// Polls with a bare TCP connect, not a full `IggyBridge::connect`: the latter carries /// `RECONNECTION_RETRIES` (3 retries after the first attempt, 4 dials at ~1s apart total) on /// every failed attempt, so a poll loop built diff --git a/gateways/kafka/tests/produce_real_bridge_tests.rs b/gateways/kafka/tests/produce_real_bridge_tests.rs index 1db4ee69f..07d4fa81d 100644 --- a/gateways/kafka/tests/produce_real_bridge_tests.rs +++ b/gateways/kafka/tests/produce_real_bridge_tests.rs @@ -651,7 +651,8 @@ async fn given_a_stopped_server_when_handled_should_answer_retriable() { let data_dir = tempfile::tempdir().expect("tempdir"); let server = TestServer::spawn(data_dir.path()).await; let state = gateway_with_topic(&server, 1).await; - drop(server); + // The Produce session connects at its first send, after this, so the port stays ours. + let _port = server.stop(); let entries = [(0, batch(&[record(0, None, b"v", &[])], Compression::None))]; let [(_, code, _)] = produce(&state, 3, 1, TOPIC, &entries).await[..] else {
