hubcio commented on code in PR #4502:
URL: https://github.com/apache/iggy/pull/4502#discussion_r4255878952
##########
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:
critical: the pool can evict a session while its detached SDK write still
runs, so a replacement connection can reorder the producer's batches. keep the
session until that write finishes, or close its transport before eviction.
##########
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:
warning: `decode_cost` skips header slots that `scan_records` charges, so
records with many headers bypass the memory limits. include header slots in the
admission cost before decoding.
##########
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:
warning: the existing timeout path releases memory permits while detached
sends still retain batches, so more producers can exceed the memory cap. keep
each reservation until all sends that own its buffers finish.
--
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]
