krishvishal commented on code in PR #4502:
URL: https://github.com/apache/iggy/pull/4502#discussion_r4259824232
##########
gateways/kafka/src/bridge/iggy_bridge/produce.rs:
##########
@@ -19,58 +19,171 @@
//!
//! 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
Review Comment:
Fixed, a session is closed only when nothing runs on it, and a send we give
up on closes its connection before the producer can send again.
##########
gateways/kafka/src/protocol/handlers/produce.rs:
##########
@@ -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);
Review Comment:
Fixed, decode_cost now counts header slots the same way scan_records does.
##########
gateways/kafka/src/protocol/handlers/produce.rs:
##########
@@ -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 {
Review Comment:
Fixed, the memory permits now stay with the send until it ends.
##########
gateways/kafka/src/protocol/handlers/produce.rs:
##########
@@ -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 {
Review Comment:
Yes, did that: both are Arc now and the owned permits move with the send.
##########
gateways/kafka/src/protocol/handlers/produce.rs:
##########
@@ -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 {
Review Comment:
Done, the producer's turn is taken before admit now, so a waiting retry
holds no slot and no memory.
##########
gateways/kafka/src/protocol/handlers/produce.rs:
##########
@@ -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)
Review Comment:
Dropped the ordering claim for these in the README, and
IGGY_KAFKA_PRODUCE_SESSIONS=0 gives the old one client behaviour if someone
needs it.
--
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]
