use std::{collections::VecDeque, time::Duration};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpStream,
sync::{mpsc, oneshot},
time,
};
use crate::native::{
KafkaClientError, KafkaClientResult,
protocol::{Decoder, Encoder},
};
const DEFAULT_COMMAND_BUFFER: usize = 64;
const MAX_RESPONSE_BYTES: usize = 512 * 1024 * 1024;
#[derive(Debug, Clone)]
pub(crate) struct BrokerConnection {
sender: mpsc::Sender<Command>,
}
#[derive(Debug)]
struct Command {
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
respond: oneshot::Sender<KafkaClientResult<Vec<u8>>>,
}
#[derive(Debug)]
struct Pending {
correlation_id: i32,
response_header_version: i16,
respond: oneshot::Sender<KafkaClientResult<Vec<u8>>>,
}
impl BrokerConnection {
pub(crate) async fn connect(
addr: String,
client_id: String,
max_in_flight: usize,
request_timeout: Duration,
) -> KafkaClientResult<Self> {
let stream = TcpStream::connect(&addr).await?;
stream.set_nodelay(true)?;
let (sender, receiver) = mpsc::channel(DEFAULT_COMMAND_BUFFER);
tokio::spawn(connection_task(
stream,
receiver,
client_id,
max_in_flight.max(1),
request_timeout,
));
Ok(Self { sender })
}
pub(crate) async fn request(
&self,
api_key: i16,
api_version: i16,
request_header_version: i16,
response_header_version: i16,
body: Vec<u8>,
) -> KafkaClientResult<Vec<u8>> {
let (respond, receiver) = oneshot::channel();
let command = Command {
api_key,
api_version,
request_header_version,
response_header_version,
body,
respond,
};
self.sender
.send(command)
.await
.map_err(|_| KafkaClientError::ChannelClosed)?;
receiver
.await
.map_err(|_| KafkaClientError::ChannelClosed)?
}
}
async fn connection_task(
mut stream: TcpStream,
mut receiver: mpsc::Receiver<Command>,
client_id: String,
max_in_flight: usize,
request_timeout: Duration,
) {
let mut correlation_id = 1_i32;
let mut pending = VecDeque::<Pending>::new();
let mut receiving_commands = true;
loop {
if !receiving_commands {
if pending.is_empty() {
break;
}
match read_response(&mut stream, &mut pending, request_timeout).await {
Ok(()) => continue,
Err(error) => {
fail_all(&mut pending, error);
break;
}
}
}
if pending.len() >= max_in_flight {
match read_response(&mut stream, &mut pending, request_timeout).await {
Ok(()) => continue,
Err(error) => {
fail_all(&mut pending, error);
break;
}
}
}
tokio::select! {
command = receiver.recv(), if receiving_commands => {
let Some(command) = command else {
receiving_commands = false;
continue;
};
let request_correlation = correlation_id;
correlation_id = if correlation_id == i32::MAX { 1 } else { correlation_id + 1 };
match write_request(&mut stream, &client_id, request_correlation, &command).await {
Ok(()) => {
pending.push_back(Pending {
correlation_id: request_correlation,
response_header_version: command.response_header_version,
respond: command.respond,
});
}
Err(error) => {
let _ = command.respond.send(Err(error));
fail_all(&mut pending, KafkaClientError::ChannelClosed);
break;
}
}
}
result = read_response(&mut stream, &mut pending, request_timeout), if !pending.is_empty() => {
if let Err(error) = result {
fail_all(&mut pending, error);
break;
}
}
}
}
}
async fn write_request(
stream: &mut TcpStream,
client_id: &str,
correlation_id: i32,
command: &Command,
) -> KafkaClientResult<()> {
let mut frame = Encoder::with_capacity(command.body.len() + 64);
frame.put_i16(command.api_key);
frame.put_i16(command.api_version);
frame.put_i32(correlation_id);
frame.put_nullable_string(Some(client_id))?;
if command.request_header_version >= 2 {
frame.put_empty_tags();
}
frame.put_raw(&command.body);
let frame = frame.into_inner();
let len = i32::try_from(frame.len())
.map_err(|_| KafkaClientError::protocol("Kafka request frame exceeds int32 length"))?;
stream.write_all(&len.to_be_bytes()).await?;
stream.write_all(&frame).await?;
Ok(())
}
async fn read_response(
stream: &mut TcpStream,
pending: &mut VecDeque<Pending>,
request_timeout: Duration,
) -> KafkaClientResult<()> {
let mut len_buf = [0_u8; 4];
time::timeout(request_timeout, stream.read_exact(&mut len_buf))
.await
.map_err(|_| KafkaClientError::protocol("Kafka response read timed out"))??;
let len = i32::from_be_bytes(len_buf);
if len < 4 {
return Err(KafkaClientError::protocol(format!(
"invalid Kafka response frame length {len}"
)));
}
let len = usize::try_from(len)
.map_err(|_| KafkaClientError::protocol("Kafka response length conversion failed"))?;
if len > MAX_RESPONSE_BYTES {
return Err(KafkaClientError::protocol(format!(
"Kafka response frame too large: {len} bytes"
)));
}
let mut frame = vec![0_u8; len];
time::timeout(request_timeout, stream.read_exact(&mut frame))
.await
.map_err(|_| KafkaClientError::protocol("Kafka response body read timed out"))??;
let Some(pending_request) = pending.pop_front() else {
return Err(KafkaClientError::protocol(
"Kafka response arrived with no pending request",
));
};
let mut decoder = Decoder::new(&frame);
let correlation_id = decoder.get_i32()?;
if pending_request.response_header_version >= 1 {
decoder.skip_tags()?;
}
if correlation_id != pending_request.correlation_id {
let _ = pending_request
.respond
.send(Err(KafkaClientError::protocol(format!(
"Kafka response correlation mismatch: expected {}, got {correlation_id}",
pending_request.correlation_id
))));
return Err(KafkaClientError::protocol(
"Kafka connection correlation-id mismatch",
));
}
let body = decoder.take(decoder.remaining())?.to_vec();
let _ = pending_request.respond.send(Ok(body));
Ok(())
}
fn fail_all(pending: &mut VecDeque<Pending>, error: KafkaClientError) {
let message = error.to_string();
while let Some(pending) = pending.pop_front() {
let _ = pending
.respond
.send(Err(KafkaClientError::protocol(message.clone())));
}
}