1pub mod actions;
55
56pub use actions::KafkaClientProtocol;
57
58use anyhow::{anyhow, bail, Context, Result};
59use bytes::Bytes;
60use std::collections::{HashMap, VecDeque};
61use std::net::SocketAddr;
62use std::sync::Arc;
63use std::time::Duration;
64use tokio::io::{AsyncReadExt, AsyncWriteExt};
65use tokio::net::TcpStream;
66use tokio::sync::{mpsc, Mutex};
67use tracing::{debug, error, info, trace, warn};
68
69use crate::client::kafka::actions::{
70 kafka_error_name, KAFKA_CLIENT_CONNECTED_EVENT, KAFKA_CLIENT_MESSAGE_DELIVERED_EVENT,
71 KAFKA_CLIENT_METADATA_RECEIVED_EVENT, KAFKA_CLIENT_RECORDS_RECEIVED_EVENT,
72};
73use crate::client::llm_budget::call_llm_for_client;
74use crate::llm::actions::client_trait::{Client, ClientActionResult};
75use crate::llm::ollama_client::OllamaClient;
76use crate::llm::ClientLlmResult;
77use crate::protocol::{Event, StartupParams};
78use crate::state::app_state::AppState;
79use crate::state::{ClientId, ClientStatus};
80
81use crate::server::kafka::kafka_protocol;
84use crate::server::kafka::{decode_field, encode_field};
85
86use kafka_protocol::messages::{
87 ApiKey, ApiVersionsRequest, ApiVersionsResponse, FetchRequest, FetchResponse, MetadataRequest,
88 MetadataResponse, OffsetCommitRequest, OffsetCommitResponse, ProduceRequest, ProduceResponse,
89 RequestHeader, ResponseHeader,
90};
91use kafka_protocol::protocol::{Decodable, Encodable, StrBytes};
92use kafka_protocol::records::{
93 Compression, Record, RecordBatchDecoder, RecordBatchEncoder, RecordEncodeOptions, TimestampType,
94};
95
96const IO_TIMEOUT: Duration = Duration::from_secs(30);
99
100const MAX_RESPONSE_BYTES: i64 = 100 * 1024 * 1024;
103
104const CLIENT_MAX_API_VERSIONS: i16 = 3;
105const CLIENT_MAX_METADATA: i16 = 12;
106const CLIENT_MAX_PRODUCE: i16 = 11;
107const CLIENT_MAX_FETCH: i16 = 16;
108const CLIENT_MAX_OFFSET_COMMIT: i16 = 9;
109const FETCH_BY_TOPIC_ID: i16 = 13;
111const METADATA_WITH_TOPIC_IDS: i16 = 10;
113const MAX_TOPIC_IDS: usize = 1024;
115
116const MAX_EVENT_RECORDS: usize = 20;
119const MAX_EVENT_VALUE_BYTES: usize = 1024;
120
121const DEFAULT_PARTITION_MAX_BYTES: i32 = 1024 * 1024;
122const DEFAULT_POLL_INTERVAL_MS: u64 = 1000;
123const MIN_POLL_INTERVAL_MS: u64 = 50;
124const DEFAULT_CLIENT_ID: &str = "netget-kafka-client";
125const DEFAULT_GROUP_ID: &str = "netget-consumer-group";
126
127type Compressor = fn(&mut bytes::BytesMut, &mut Vec<u8>, Compression) -> Result<()>;
130type Decompressor = fn(&mut Bytes, Compression) -> Result<std::io::Cursor<Bytes>>;
131
132#[derive(Clone, Debug)]
134struct Assignment {
135 topic: String,
136 partition: i32,
137 next_offset: i64,
138}
139
140struct KafkaConn {
142 stream: TcpStream,
143 next_correlation_id: i32,
144 versions: HashMap<i16, (i16, i16)>,
146 client_id: StrBytes,
147 last_request_bytes: usize,
152 topic_ids: HashMap<String, uuid::Uuid>,
154}
155
156impl KafkaConn {
157 async fn topic_id(&mut self, topic: &str) -> Result<Option<uuid::Uuid>> {
161 if let Some(id) = self.topic_ids.get(topic) {
162 return Ok(Some(*id));
163 }
164 use kafka_protocol::messages::metadata_request::MetadataRequestTopic;
165 let version = self.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?;
166 if version < METADATA_WITH_TOPIC_IDS {
167 return Ok(None);
168 }
169 let response: MetadataResponse = self
170 .exchange(
171 ApiKey::Metadata,
172 version,
173 &MetadataRequest::default()
174 .with_topics(Some(vec![MetadataRequestTopic::default()
175 .with_name(Some(StrBytes::from_string(topic.to_string()).into()))])),
176 )
177 .await
178 .context("Metadata request for a topic id failed")?;
179 self.remember_topic_ids(&response);
180 Ok(self.topic_ids.get(topic).copied())
181 }
182
183 fn remember_topic_ids(&mut self, response: &MetadataResponse) {
185 for t in &response.topics {
186 let Some(name) = t.name.as_ref() else {
187 continue;
188 };
189 if t.error_code != 0 || t.topic_id.is_nil() {
190 continue;
191 }
192 if self.topic_ids.len() >= MAX_TOPIC_IDS && !self.topic_ids.contains_key(name.as_str())
193 {
194 self.topic_ids.clear();
195 }
196 self.topic_ids.insert(name.to_string(), t.topic_id);
197 }
198 }
199
200 fn version_for(&self, api_key: ApiKey, our_max: i16) -> Result<i16> {
206 let (broker_min, broker_max) =
207 *self.versions.get(&(api_key as i16)).with_context(|| {
208 format!(
209 "broker does not advertise {api_key:?} (key {})",
210 api_key as i16
211 )
212 })?;
213 let chosen = our_max.min(broker_max);
214 if chosen < broker_min {
215 bail!(
216 "broker supports {api_key:?} v{broker_min}-v{broker_max}, but this client \
217 implements at most v{our_max}"
218 );
219 }
220 Ok(chosen)
221 }
222
223 fn build_request<B: Encodable>(
224 &mut self,
225 api_key: ApiKey,
226 api_version: i16,
227 body: &B,
228 ) -> Result<(i32, Vec<u8>)> {
229 let correlation_id = self.next_correlation_id;
230 self.next_correlation_id = self.next_correlation_id.wrapping_add(1);
231
232 let mut buf = Vec::new();
233 RequestHeader::default()
234 .with_request_api_key(api_key as i16)
235 .with_request_api_version(api_version)
236 .with_correlation_id(correlation_id)
237 .with_client_id(Some(self.client_id.clone()))
238 .encode(&mut buf, api_key.request_header_version(api_version))
239 .with_context(|| format!("encoding {api_key:?} v{api_version} request header"))?;
240 body.encode(&mut buf, api_version)
241 .with_context(|| format!("encoding {api_key:?} v{api_version} request body"))?;
242 Ok((correlation_id, buf))
243 }
244
245 async fn write_frame(&mut self, body: &[u8]) -> Result<()> {
246 let size = i32::try_from(body.len()).context("request does not fit in an i32 frame")?;
247 tokio::time::timeout(IO_TIMEOUT, self.stream.write_all(&size.to_be_bytes()))
248 .await
249 .context("timed out writing Kafka frame size")??;
250 tokio::time::timeout(IO_TIMEOUT, self.stream.write_all(body))
251 .await
252 .context("timed out writing Kafka frame body")??;
253 self.last_request_bytes = 4 + body.len();
254 Ok(())
255 }
256
257 async fn read_frame(&mut self) -> Result<Vec<u8>> {
258 let mut size = [0u8; 4];
259 tokio::time::timeout(IO_TIMEOUT, self.stream.read_exact(&mut size))
260 .await
261 .context("timed out reading Kafka response size")??;
262 let announced = i64::from(i32::from_be_bytes(size));
264 if announced <= 0 || announced > MAX_RESPONSE_BYTES {
265 bail!("broker announced an implausible response size of {announced} bytes");
266 }
267 let mut buf = vec![0u8; announced as usize];
268 tokio::time::timeout(IO_TIMEOUT, self.stream.read_exact(&mut buf))
269 .await
270 .context("timed out reading Kafka response body")??;
271 Ok(buf)
272 }
273
274 async fn exchange<B: Encodable, R: Decodable>(
276 &mut self,
277 api_key: ApiKey,
278 api_version: i16,
279 body: &B,
280 ) -> Result<R> {
281 let (correlation_id, request) = self.build_request(api_key, api_version, body)?;
282 self.write_frame(&request).await?;
283 let response = self.read_frame().await?;
284
285 let mut cursor = std::io::Cursor::new(response.as_slice());
286 let header =
287 ResponseHeader::decode(&mut cursor, api_key.response_header_version(api_version))
288 .with_context(|| format!("decoding {api_key:?} v{api_version} response header"))?;
289 if header.correlation_id != correlation_id {
290 bail!(
291 "broker answered {api_key:?} with correlation id {} instead of {correlation_id}; \
292 the connection is desynchronised",
293 header.correlation_id
294 );
295 }
296 R::decode(&mut cursor, api_version)
297 .with_context(|| format!("decoding {api_key:?} v{api_version} response body"))
298 }
299
300 async fn send_only<B: Encodable>(
303 &mut self,
304 api_key: ApiKey,
305 api_version: i16,
306 body: &B,
307 ) -> Result<()> {
308 let (_, request) = self.build_request(api_key, api_version, body)?;
309 self.write_frame(&request).await
310 }
311}
312
313#[derive(Clone)]
315struct Session {
316 conn: Arc<Mutex<KafkaConn>>,
317 assignments: Arc<Mutex<Vec<Assignment>>>,
318 memory: Arc<Mutex<String>>,
319 client_id: ClientId,
320 app_state: Arc<AppState>,
321 llm_client: OllamaClient,
322 status_tx: mpsc::UnboundedSender<String>,
323 default_group_id: String,
324 remote_addr: String,
325}
326
327pub struct KafkaClient;
329
330impl KafkaClient {
331 pub async fn connect_with_llm_actions(
333 remote_addr: String,
334 llm_client: OllamaClient,
335 app_state: Arc<AppState>,
336 status_tx: mpsc::UnboundedSender<String>,
337 client_id: ClientId,
338 startup_params: Option<StartupParams>,
339 ) -> Result<SocketAddr> {
340 let params = startup_params.as_ref();
343 let client_id_str = params
344 .map(|p| p.get_optional_string("client_id"))
345 .transpose()?
346 .flatten()
347 .unwrap_or_else(|| DEFAULT_CLIENT_ID.to_string());
348
349 let topics: Vec<String> = params
350 .map(|p| p.get_optional_array("topics"))
351 .transpose()?
352 .flatten()
353 .map(|arr| {
354 arr.iter()
355 .filter_map(|v| v.as_str())
356 .filter(|s| !s.is_empty())
357 .map(|s| s.to_string())
358 .collect()
359 })
360 .unwrap_or_default();
361
362 let partition_raw = params
363 .map(|p| p.get_optional_i64("partition"))
364 .transpose()?
365 .flatten()
366 .unwrap_or(0);
367 if !(0..=actions::MAX_PARTITION).contains(&partition_raw) {
368 bail!(
369 "Kafka client 'partition' must be between 0 and {}, got {partition_raw}",
370 actions::MAX_PARTITION
371 );
372 }
373 let partition = partition_raw as i32;
374
375 let start_offset = params
376 .map(|p| p.get_optional_i64("start_offset"))
377 .transpose()?
378 .flatten()
379 .unwrap_or(0);
380 if start_offset < 0 {
381 bail!("Kafka client 'start_offset' must not be negative, got {start_offset}");
382 }
383
384 let poll_interval_raw = params
385 .map(|p| p.get_optional_u64("poll_interval_ms"))
386 .transpose()?
387 .flatten()
388 .unwrap_or(DEFAULT_POLL_INTERVAL_MS);
389 if poll_interval_raw < MIN_POLL_INTERVAL_MS {
390 bail!(
391 "Kafka client 'poll_interval_ms' must be at least {MIN_POLL_INTERVAL_MS}, got \
392 {poll_interval_raw}"
393 );
394 }
395 let poll_interval = Duration::from_millis(poll_interval_raw);
396
397 let group_id = params
398 .map(|p| p.get_optional_string("group_id"))
399 .transpose()?
400 .flatten()
401 .unwrap_or_else(|| DEFAULT_GROUP_ID.to_string());
402
403 info!(
404 "Kafka client {} connecting to {} (client_id={}, topics={:?})",
405 client_id, remote_addr, client_id_str, topics
406 );
407
408 let stream = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect(&remote_addr))
409 .await
410 .with_context(|| format!("timed out connecting to Kafka broker {remote_addr}"))?
411 .with_context(|| format!("failed to connect to Kafka broker {remote_addr}"))?;
412 let local_addr = stream.local_addr()?;
413
414 let mut conn = KafkaConn {
415 stream,
416 next_correlation_id: 1,
417 versions: HashMap::new(),
418 client_id: StrBytes::from_string(client_id_str.clone()),
419 last_request_bytes: 0,
420 topic_ids: HashMap::new(),
421 };
422
423 conn.versions = Self::negotiate_api_versions(&mut conn).await?;
425
426 let negotiated = serde_json::json!({
427 "metadata": conn.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?,
428 "produce": conn.version_for(ApiKey::Produce, CLIENT_MAX_PRODUCE)?,
429 "fetch": conn.version_for(ApiKey::Fetch, CLIENT_MAX_FETCH)?,
430 "offset_commit": conn.version_for(ApiKey::OffsetCommit, CLIENT_MAX_OFFSET_COMMIT)?,
431 });
432
433 let requested: Option<Vec<String>> = if topics.is_empty() {
435 None
436 } else {
437 Some(topics.clone())
438 };
439 let mut metadata = Self::request_metadata(&mut conn, requested.as_deref()).await?;
440
441 info!(
442 "Kafka client {} negotiated {} and read metadata for {} topic(s)",
443 client_id,
444 negotiated,
445 metadata
446 .get("topics")
447 .and_then(|t| t.as_array())
448 .map(|a| a.len())
449 .unwrap_or(0)
450 );
451
452 app_state
453 .update_client_status(client_id, ClientStatus::Connected)
454 .await;
455 let _ = status_tx.send(format!(
456 "[CLIENT] Kafka client {client_id} connected to {remote_addr}"
457 ));
458 let _ = status_tx.send("__UPDATE_UI__".to_string());
459
460 app_state
461 .with_client_mut(client_id, |client_inst| {
462 client_inst
463 .set_protocol_field("brokers".to_string(), serde_json::json!(remote_addr));
464 client_inst.set_protocol_field(
465 "kafka_client_id".to_string(),
466 serde_json::json!(client_id_str),
467 );
468 client_inst.set_protocol_field("api_versions".to_string(), negotiated.clone());
469 })
470 .await;
471
472 let assignments: Vec<Assignment> = topics
473 .iter()
474 .map(|topic| Assignment {
475 topic: topic.clone(),
476 partition,
477 next_offset: start_offset,
478 })
479 .collect();
480 let polling = !assignments.is_empty();
481
482 let memory = app_state
483 .get_memory_for_client(client_id)
484 .await
485 .unwrap_or_default();
486
487 let session = Session {
488 conn: Arc::new(Mutex::new(conn)),
489 assignments: Arc::new(Mutex::new(assignments)),
490 memory: Arc::new(Mutex::new(memory)),
491 client_id,
492 app_state: app_state.clone(),
493 llm_client,
494 status_tx: status_tx.clone(),
495 default_group_id: group_id,
496 remote_addr: remote_addr.clone(),
497 };
498
499 let connected_data = {
502 let obj = metadata
503 .as_object_mut()
504 .expect("request_metadata builds an object");
505 obj.insert("remote_addr".to_string(), serde_json::json!(remote_addr));
506 obj.insert("api_versions".to_string(), negotiated);
507 metadata
508 };
509
510 let command_rx =
515 crate::client::command_support::register_command_channel(&app_state, client_id).await;
516 let command_session = session.clone();
517 let command_task = tokio::spawn(async move {
518 command_session.command_loop(command_rx).await;
519 });
520 app_state
521 .register_client_task(client_id, command_task)
522 .await;
523
524 let task_registrar = app_state.clone();
525 let task_handle = tokio::spawn(async move {
526 let stop = session
527 .drive(Event::new(&KAFKA_CLIENT_CONNECTED_EVENT, connected_data))
528 .await;
529
530 if stop {
531 session.close("handler asked to disconnect").await;
532 return;
533 }
534 if !polling {
535 debug!(
536 "Kafka client {} has no topics to poll; the connection stays open for \
537 scheduled tasks and further instructions",
538 session.client_id
539 );
540 return;
541 }
542 session.poll_loop(poll_interval).await;
543 });
544 task_registrar
545 .register_client_task(client_id, task_handle)
546 .await;
547
548 Ok(local_addr)
549 }
550
551 async fn negotiate_api_versions(conn: &mut KafkaConn) -> Result<HashMap<i16, (i16, i16)>> {
559 let request = ApiVersionsRequest::default()
560 .with_client_software_name(StrBytes::from_static_str("netget"))
561 .with_client_software_version(StrBytes::from_static_str(env!("CARGO_PKG_VERSION")));
562
563 let (correlation_id, bytes) =
564 conn.build_request(ApiKey::ApiVersions, CLIENT_MAX_API_VERSIONS, &request)?;
565 conn.write_frame(&bytes).await?;
566 let response = conn.read_frame().await?;
567
568 let decoded = Self::decode_api_versions(&response, CLIENT_MAX_API_VERSIONS, correlation_id)
569 .and_then(|(code, body)| {
570 if code == 0 {
571 Ok(body)
572 } else {
573 Err(anyhow!(
574 "ApiVersions v{CLIENT_MAX_API_VERSIONS} refused with error {code}"
575 ))
576 }
577 });
578
579 let body = match decoded {
580 Ok(body) => body,
581 Err(first) => {
582 debug!("Kafka client stepping ApiVersions down to v0: {}", first);
583 let (code, body) = Self::decode_api_versions(&response, 0, correlation_id)
584 .context("broker's ApiVersions reply is unreadable at v3 and at v0")?;
585 if code != 0 && body.api_keys.is_empty() {
586 bail!(
587 "broker refused ApiVersions with error {code} ({}) and sent no supported-API \
588 table, so there is no version to step down to",
589 kafka_error_name(code)
590 );
591 }
592 body
593 }
594 };
595
596 if body.api_keys.is_empty() {
597 bail!("broker advertised no APIs at all");
598 }
599
600 Ok(body
601 .api_keys
602 .iter()
603 .map(|a| (a.api_key, (a.min_version, a.max_version)))
604 .collect())
605 }
606
607 fn decode_api_versions(
608 bytes: &[u8],
609 api_version: i16,
610 expected_correlation_id: i32,
611 ) -> Result<(i16, ApiVersionsResponse)> {
612 let mut cursor = std::io::Cursor::new(bytes);
613 let header = ResponseHeader::decode(
614 &mut cursor,
615 ApiKey::ApiVersions.response_header_version(api_version),
616 )
617 .context("decoding ApiVersions response header")?;
618 if header.correlation_id != expected_correlation_id {
619 bail!(
620 "broker answered ApiVersions with correlation id {} instead of {}",
621 header.correlation_id,
622 expected_correlation_id
623 );
624 }
625 let body = ApiVersionsResponse::decode(&mut cursor, api_version)
626 .context("decoding ApiVersions response body")?;
627 Ok((body.error_code, body))
628 }
629
630 async fn request_metadata(
632 conn: &mut KafkaConn,
633 topics: Option<&[String]>,
634 ) -> Result<serde_json::Value> {
635 use kafka_protocol::messages::metadata_request::MetadataRequestTopic;
636
637 let version = conn.version_for(ApiKey::Metadata, CLIENT_MAX_METADATA)?;
638
639 let named: Option<Vec<MetadataRequestTopic>> = match topics {
640 Some(names) => Some(
641 names
642 .iter()
643 .map(|n| {
644 MetadataRequestTopic::default()
645 .with_name(Some(StrBytes::from_string(n.clone()).into()))
646 })
647 .collect(),
648 ),
649 None if version == 0 => Some(Vec::new()),
651 None => None,
652 };
653
654 let response: MetadataResponse = conn
655 .exchange(
656 ApiKey::Metadata,
657 version,
658 &MetadataRequest::default().with_topics(named),
659 )
660 .await
661 .context("Metadata request failed")?;
662 conn.remember_topic_ids(&response);
665
666 Ok(metadata_to_json(&response))
667 }
668}
669
670impl Session {
671 async fn drive(&self, first: Event) -> bool {
678 let Some(instruction) = self
679 .app_state
680 .get_instruction_for_client(self.client_id)
681 .await
682 else {
683 return false;
684 };
685
686 let mut queue: VecDeque<Event> = VecDeque::from([first]);
687 while let Some(event) = queue.pop_front() {
688 let memory = self.memory.lock().await.clone();
689 let protocol = KafkaClientProtocol::new();
690
691 let result = call_llm_for_client(
692 &self.llm_client,
693 &self.app_state,
694 self.client_id.to_string(),
695 &instruction,
696 &memory,
697 Some(&event),
698 &protocol,
699 &self.status_tx,
700 )
701 .await;
702
703 let ClientLlmResult {
704 actions,
705 memory_updates,
706 } = match result {
707 Ok(r) => r,
708 Err(e) => {
709 error!(
710 "Kafka client {} could not handle {}: {}",
711 self.client_id,
712 event.id(),
713 e
714 );
715 let _ = self
716 .status_tx
717 .send(format!("[CLIENT] Kafka handler failed: {e}"));
718 return false;
719 }
720 };
721 if let Some(mem) = memory_updates {
722 *self.memory.lock().await = mem;
723 }
724
725 for action in actions {
726 match protocol.execute_action(action) {
727 Ok(ClientActionResult::Disconnect) => return true,
728 Ok(ClientActionResult::WaitForMore) | Ok(ClientActionResult::NoAction) => {}
729 Ok(ClientActionResult::Custom { name, data }) => {
730 match self.perform(&name, &data).await {
731 Ok((Some(follow_up), _)) => queue.push_back(follow_up),
732 Ok((None, _)) => {}
733 Err(e) => {
734 error!(
737 "Kafka client {} action '{}' failed: {}",
738 self.client_id, name, e
739 );
740 let _ = self
741 .status_tx
742 .send(format!("[CLIENT] Kafka {name} failed: {e}"));
743 if is_fatal(&e) {
744 self.mark_error(&e.to_string()).await;
745 return true;
746 }
747 }
748 }
749 }
750 Ok(other) => {
751 warn!(
752 "Kafka client {} ignoring unsupported action result {:?}",
753 self.client_id, other
754 );
755 }
756 Err(e) => {
757 error!("Kafka client {} rejected an action: {}", self.client_id, e);
758 let _ = self
759 .status_tx
760 .send(format!("[CLIENT] Kafka action rejected: {e}"));
761 }
762 }
763 }
764 }
765 false
766 }
767
768 async fn command_loop(
782 &self,
783 mut command_rx: mpsc::Receiver<crate::state::client_handles::ClientCommand>,
784 ) {
785 use crate::llm::actions::protocol_trait::Protocol;
786 use crate::state::client_handles::ClientSendOutcome;
787 use crate::state::AccessLogOwner;
788
789 let protocol = KafkaClientProtocol::new();
790
791 while let Some(command) = command_rx.recv().await {
792 let action = command.action.clone();
793 let mut follow_up: Option<Event> = None;
794 let mut disconnect = false;
795
796 let outcome: Result<ClientSendOutcome> = match protocol.execute_action(action.clone()) {
797 Err(e) => Ok(ClientSendOutcome::Rejected {
798 error: e.to_string(),
799 }),
800 Ok(ClientActionResult::Disconnect) => {
801 disconnect = true;
802 Ok(ClientSendOutcome::Disconnected)
803 }
804 Ok(ClientActionResult::WaitForMore) => Ok(ClientSendOutcome::Executed {
805 detail: "wait_for_more: nothing was asked of the broker".to_string(),
806 }),
807 Ok(ClientActionResult::NoAction) => Ok(ClientSendOutcome::Executed {
808 detail: "no_action".to_string(),
809 }),
810 Ok(ClientActionResult::Custom { name, data }) => {
811 match self.perform(&name, &data).await {
812 Ok((event, wrote)) => {
813 follow_up = event;
814 Ok(ClientSendOutcome::Sent { bytes_sent: wrote })
815 }
816 Err(e) => {
817 if is_fatal(&e) {
820 self.mark_error(&e.to_string()).await;
821 }
822 Err(e.context(format!("injected Kafka action '{name}'")))
823 }
824 }
825 }
826 Ok(other) => Ok(ClientSendOutcome::Executed {
827 detail: format!("unsupported action result {other:?}"),
828 }),
829 };
830
831 let outcome_json = match &outcome {
832 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
833 Err(e) => serde_json::json!({"error": e.to_string()}),
834 };
835 self.app_state
836 .record_access_log(
837 AccessLogOwner::Client(self.client_id.as_u32()),
838 protocol.protocol_name(),
839 None,
840 "injected_action",
841 action,
842 vec![outcome_json],
843 )
844 .await;
845
846 if let Err(e) = &outcome {
847 error!(
848 "Kafka client {} injected action failed: {}",
849 self.client_id, e
850 );
851 let _ = self
852 .status_tx
853 .send(format!("[CLIENT] Kafka injected action failed: {e}"));
854 }
855 let _ = self.status_tx.send("__UPDATE_UI__".to_string());
856 crate::client::command_support::reply(command, outcome);
857
858 if disconnect {
859 self.close("injected disconnect").await;
860 break;
861 }
862
863 if let Some(event) = follow_up {
864 let session = self.clone();
865 let handle = tokio::spawn(async move {
866 if session.drive(event).await {
867 session.close("handler asked to disconnect").await;
868 }
869 });
870 self.app_state
871 .register_client_task(self.client_id, handle)
872 .await;
873 }
874 }
875
876 self.app_state.remove_client_handle(self.client_id).await;
879 }
880
881 async fn perform(
889 &self,
890 name: &str,
891 data: &serde_json::Value,
892 ) -> Result<(Option<Event>, usize)> {
893 match name {
894 "kafka_produce" => self.produce(data).await,
895 "kafka_fetch" => {
896 let topic = str_field(data, "topic")?;
897 let partition = i64_field(data, "partition")? as i32;
898 let offset = match data.get("offset").and_then(|v| v.as_i64()) {
899 Some(o) => o,
900 None => self.tracked_offset(&topic, partition).await.unwrap_or(0),
901 };
902 let max_bytes = data
903 .get("max_bytes")
904 .and_then(|v| v.as_i64())
905 .unwrap_or(i64::from(DEFAULT_PARTITION_MAX_BYTES))
906 as i32;
907 self.fetch(&topic, partition, offset, max_bytes).await
908 }
909 "kafka_metadata" => {
910 let topics: Option<Vec<String>> =
911 data.get("topics").and_then(|v| v.as_array()).map(|a| {
912 a.iter()
913 .filter_map(|v| v.as_str())
914 .map(|s| s.to_string())
915 .collect()
916 });
917 let mut conn = self.conn.lock().await;
918 let metadata = KafkaClient::request_metadata(&mut conn, topics.as_deref()).await?;
919 let wrote = conn.last_request_bytes;
920 drop(conn);
921 Ok((
922 Some(Event::new(&KAFKA_CLIENT_METADATA_RECEIVED_EVENT, metadata)),
923 wrote,
924 ))
925 }
926 "kafka_commit" => self.commit(data).await.map(|wrote| (None, wrote)),
927 other => Err(anyhow!("unknown Kafka client action result '{other}'")),
928 }
929 }
930
931 async fn produce(&self, data: &serde_json::Value) -> Result<(Option<Event>, usize)> {
933 use kafka_protocol::messages::produce_request::{PartitionProduceData, TopicProduceData};
934
935 let topic = str_field(data, "topic")?;
936 let partition = i64_field(data, "partition")? as i32;
937 let acks = i64_field(data, "acks")? as i16;
938
939 let key = decode_field(
943 data.get("key"),
944 data.get("key_encoding")
945 .and_then(|v| v.as_str())
946 .unwrap_or("utf8"),
947 )
948 .context("record key")?;
949 let value = decode_field(
950 data.get("value"),
951 data.get("value_encoding")
952 .and_then(|v| v.as_str())
953 .unwrap_or("utf8"),
954 )
955 .context("record value")?;
956
957 let batch = encode_single_record(key, value)?;
958
959 let request = ProduceRequest::default()
960 .with_acks(acks)
961 .with_timeout_ms(IO_TIMEOUT.as_millis() as i32)
962 .with_topic_data(vec![TopicProduceData::default()
963 .with_name(StrBytes::from_string(topic.clone()).into())
964 .with_partition_data(vec![PartitionProduceData::default()
965 .with_index(partition)
966 .with_records(Some(batch))])]);
967
968 let mut conn = self.conn.lock().await;
969 let version = conn.version_for(ApiKey::Produce, CLIENT_MAX_PRODUCE)?;
970
971 if acks == 0 {
972 conn.send_only(ApiKey::Produce, version, &request).await?;
975 let wrote = conn.last_request_bytes;
976 drop(conn);
977 info!(
978 "Kafka client {} produced to {}/{} with acks=0 (no acknowledgement requested)",
979 self.client_id, topic, partition
980 );
981 let _ = self.status_tx.send(format!(
982 "[CLIENT] Kafka produced to {topic}/{partition} (acks=0, unacknowledged)"
983 ));
984 return Ok((None, wrote));
985 }
986
987 let response: ProduceResponse = conn.exchange(ApiKey::Produce, version, &request).await?;
988 let wrote = conn.last_request_bytes;
989 drop(conn);
990
991 let topic_response = response
992 .responses
993 .iter()
994 .find(|t| t.name.as_str() == topic)
995 .or_else(|| response.responses.first())
996 .context("Produce response named no topic")?;
997 let partition_response = topic_response
998 .partition_responses
999 .iter()
1000 .find(|p| p.index == partition)
1001 .or_else(|| topic_response.partition_responses.first())
1002 .context("Produce response named no partition")?;
1003
1004 let error_code = partition_response.error_code;
1005 let delivered = error_code == 0;
1006 if delivered {
1007 info!(
1008 "Kafka client {} produced to {}/{} at offset {}",
1009 self.client_id, topic, partition, partition_response.base_offset
1010 );
1011 let _ = self.status_tx.send(format!(
1012 "[CLIENT] Kafka produced to {topic}/{partition} at offset {}",
1013 partition_response.base_offset
1014 ));
1015 } else {
1016 error!(
1018 "Kafka client {} produce to {}/{} rejected: {} ({})",
1019 self.client_id,
1020 topic,
1021 partition,
1022 error_code,
1023 kafka_error_name(error_code)
1024 );
1025 let _ = self.status_tx.send(format!(
1026 "[CLIENT] Kafka produce to {topic}/{partition} rejected: {} ({error_code})",
1027 kafka_error_name(error_code)
1028 ));
1029 }
1030
1031 Ok((
1032 Some(Event::new(
1033 &KAFKA_CLIENT_MESSAGE_DELIVERED_EVENT,
1034 serde_json::json!({
1035 "topic": topic,
1036 "partition": partition,
1037 "base_offset": partition_response.base_offset,
1038 "error_code": error_code,
1039 "error_name": kafka_error_name(error_code),
1040 "delivered": delivered,
1041 "api_version": version,
1042 }),
1043 )),
1044 wrote,
1045 ))
1046 }
1047
1048 async fn fetch(
1051 &self,
1052 topic: &str,
1053 partition: i32,
1054 offset: i64,
1055 max_bytes: i32,
1056 ) -> Result<(Option<Event>, usize)> {
1057 use kafka_protocol::messages::fetch_request::{FetchPartition, FetchTopic};
1058
1059 let mut conn = self.conn.lock().await;
1060 let mut version = conn.version_for(ApiKey::Fetch, CLIENT_MAX_FETCH)?;
1061 let mut fetch_topic = FetchTopic::default();
1064 if version >= FETCH_BY_TOPIC_ID {
1065 match conn.topic_id(topic).await? {
1066 Some(id) => fetch_topic = fetch_topic.with_topic_id(id),
1067 None => version = conn.version_for(ApiKey::Fetch, FETCH_BY_TOPIC_ID - 1)?,
1068 }
1069 }
1070 if version < FETCH_BY_TOPIC_ID {
1071 fetch_topic = fetch_topic.with_topic(StrBytes::from_string(topic.to_string()).into());
1072 }
1073 let request = FetchRequest::default()
1074 .with_max_wait_ms(500)
1075 .with_min_bytes(1)
1076 .with_max_bytes(max_bytes)
1077 .with_topics(vec![fetch_topic.with_partitions(vec![
1078 FetchPartition::default()
1079 .with_partition(partition)
1080 .with_fetch_offset(offset)
1081 .with_partition_max_bytes(max_bytes),
1082 ])]);
1083 let response: FetchResponse = conn.exchange(ApiKey::Fetch, version, &request).await?;
1084 let wrote = conn.last_request_bytes;
1085 drop(conn);
1086
1087 if response.error_code != 0 {
1088 bail!(
1089 "Fetch failed at the response level: {} ({})",
1090 response.error_code,
1091 kafka_error_name(response.error_code)
1092 );
1093 }
1094
1095 let Some(topic_response) = response.responses.first() else {
1096 debug!(
1097 "Kafka client {} fetched {}/{}: broker returned no topic",
1098 self.client_id, topic, partition
1099 );
1100 return Ok((None, wrote));
1101 };
1102 let Some(partition_data) = topic_response.partitions.first() else {
1103 return Ok((None, wrote));
1104 };
1105
1106 if partition_data.error_code != 0 {
1107 bail!(
1108 "Fetch of {topic}/{partition} failed: {} ({})",
1109 partition_data.error_code,
1110 kafka_error_name(partition_data.error_code)
1111 );
1112 }
1113
1114 let records = match partition_data.records.as_ref() {
1115 Some(raw) if !raw.is_empty() => {
1116 let owned = Bytes::copy_from_slice(raw.as_ref());
1117 let mut cursor = std::io::Cursor::new(owned);
1118 RecordBatchDecoder::decode_with_custom_compression::<_, Decompressor>(
1119 &mut cursor,
1120 None,
1121 )
1122 .context("broker's record batch did not decode")?
1123 }
1124 _ => Vec::new(),
1125 };
1126
1127 if records.is_empty() {
1128 trace!(
1129 "Kafka client {} fetched {}/{} from offset {}: nothing new",
1130 self.client_id,
1131 topic,
1132 partition,
1133 offset
1134 );
1135 return Ok((None, wrote));
1136 }
1137
1138 let last_offset = records.iter().map(|r| r.offset).max().unwrap_or(offset);
1139 let next_offset = last_offset + 1;
1140 self.set_tracked_offset(topic, partition, next_offset).await;
1141
1142 info!(
1143 "Kafka client {} fetched {} record(s) from {}/{} at offset {}",
1144 self.client_id,
1145 records.len(),
1146 topic,
1147 partition,
1148 offset
1149 );
1150 let _ = self.status_tx.send(format!(
1151 "[CLIENT] Kafka received {} record(s) from {topic}/{partition}",
1152 records.len()
1153 ));
1154
1155 Ok((
1156 Some(Event::new(
1157 &KAFKA_CLIENT_RECORDS_RECEIVED_EVENT,
1158 serde_json::json!({
1159 "topic": topic,
1160 "partition": partition,
1161 "api_version": version,
1162 "high_watermark": partition_data.high_watermark,
1163 "next_offset": next_offset,
1164 "record_count": records.len(),
1165 "records": records.iter().take(MAX_EVENT_RECORDS).map(record_to_json).collect::<Vec<_>>(),
1166 }),
1167 )),
1168 wrote,
1169 ))
1170 }
1171
1172 async fn commit(&self, data: &serde_json::Value) -> Result<usize> {
1175 use kafka_protocol::messages::offset_commit_request::{
1176 OffsetCommitRequestPartition, OffsetCommitRequestTopic,
1177 };
1178
1179 let topic = str_field(data, "topic")?;
1180 let partition = i64_field(data, "partition")? as i32;
1181 let offset = i64_field(data, "offset")?;
1182 let group_id = data
1183 .get("group_id")
1184 .and_then(|v| v.as_str())
1185 .unwrap_or(&self.default_group_id)
1186 .to_string();
1187
1188 let request = OffsetCommitRequest::default()
1189 .with_group_id(StrBytes::from_string(group_id.clone()).into())
1190 .with_generation_id_or_member_epoch(-1)
1193 .with_topics(vec![OffsetCommitRequestTopic::default()
1194 .with_name(StrBytes::from_string(topic.clone()).into())
1195 .with_partitions(vec![OffsetCommitRequestPartition::default()
1196 .with_partition_index(partition)
1197 .with_committed_offset(offset)])]);
1198
1199 let mut conn = self.conn.lock().await;
1200 let version = conn.version_for(ApiKey::OffsetCommit, CLIENT_MAX_OFFSET_COMMIT)?;
1201 let response: OffsetCommitResponse = conn
1202 .exchange(ApiKey::OffsetCommit, version, &request)
1203 .await?;
1204 let wrote = conn.last_request_bytes;
1205 drop(conn);
1206
1207 let error_code = response
1208 .topics
1209 .first()
1210 .and_then(|t| t.partitions.first())
1211 .map(|p| p.error_code)
1212 .context("OffsetCommit response named no partition")?;
1213
1214 if error_code != 0 {
1215 bail!(
1216 "commit of {topic}/{partition} offset {offset} was refused: {error_code} ({})",
1217 kafka_error_name(error_code)
1218 );
1219 }
1220
1221 info!(
1222 "Kafka client {} committed {}/{} offset {} for group {}",
1223 self.client_id, topic, partition, offset, group_id
1224 );
1225 let _ = self.status_tx.send(format!(
1226 "[CLIENT] Kafka committed {topic}/{partition} offset {offset}"
1227 ));
1228 Ok(wrote)
1229 }
1230
1231 async fn poll_loop(&self, interval: Duration) {
1237 let mut first = true;
1238 loop {
1239 if !first {
1240 tokio::time::sleep(interval).await;
1241 }
1242 first = false;
1243
1244 if self.app_state.get_client(self.client_id).await.is_none() {
1245 info!("Kafka client {} stopped", self.client_id);
1246 return;
1247 }
1248
1249 let assignments = self.assignments.lock().await.clone();
1250 for assignment in assignments {
1251 match self
1252 .fetch(
1253 &assignment.topic,
1254 assignment.partition,
1255 assignment.next_offset,
1256 DEFAULT_PARTITION_MAX_BYTES,
1257 )
1258 .await
1259 {
1260 Ok((Some(event), _)) => {
1261 if self.drive(event).await {
1262 self.close("handler asked to disconnect").await;
1263 return;
1264 }
1265 }
1266 Ok((None, _)) => {}
1267 Err(e) => {
1268 error!(
1269 "Kafka client {} polling {}/{} failed: {}",
1270 self.client_id, assignment.topic, assignment.partition, e
1271 );
1272 let _ = self
1273 .status_tx
1274 .send(format!("[CLIENT] Kafka poll failed: {e}"));
1275 if is_fatal(&e) {
1276 self.mark_error(&e.to_string()).await;
1277 return;
1278 }
1279 }
1280 }
1281 }
1282 }
1283 }
1284
1285 async fn tracked_offset(&self, topic: &str, partition: i32) -> Option<i64> {
1286 self.assignments
1287 .lock()
1288 .await
1289 .iter()
1290 .find(|a| a.topic == topic && a.partition == partition)
1291 .map(|a| a.next_offset)
1292 }
1293
1294 async fn set_tracked_offset(&self, topic: &str, partition: i32, next_offset: i64) {
1297 let mut assignments = self.assignments.lock().await;
1298 match assignments
1299 .iter_mut()
1300 .find(|a| a.topic == topic && a.partition == partition)
1301 {
1302 Some(existing) => existing.next_offset = next_offset,
1303 None => assignments.push(Assignment {
1304 topic: topic.to_string(),
1305 partition,
1306 next_offset,
1307 }),
1308 }
1309 }
1310
1311 async fn mark_error(&self, message: &str) {
1312 self.app_state.remove_client_handle(self.client_id).await;
1314 self.app_state
1315 .update_client_status(self.client_id, ClientStatus::Error(message.to_string()))
1316 .await;
1317 let _ = self.status_tx.send("__UPDATE_UI__".to_string());
1318 }
1319
1320 async fn close(&self, why: &str) {
1321 info!(
1322 "Kafka client {} disconnecting from {}: {}",
1323 self.client_id, self.remote_addr, why
1324 );
1325 let _ = self.conn.lock().await.stream.shutdown().await;
1328 self.app_state.remove_client_handle(self.client_id).await;
1330 self.app_state
1331 .update_client_status(self.client_id, ClientStatus::Disconnected)
1332 .await;
1333 let _ = self.status_tx.send(format!(
1334 "[CLIENT] Kafka client {} disconnected",
1335 self.client_id
1336 ));
1337 let _ = self.status_tx.send("__UPDATE_UI__".to_string());
1338 }
1339}
1340
1341fn is_fatal(e: &anyhow::Error) -> bool {
1346 let text = e.to_string();
1347 e.downcast_ref::<std::io::Error>().is_some()
1348 || e.chain()
1349 .any(|c| c.downcast_ref::<std::io::Error>().is_some())
1350 || text.contains("desynchronised")
1351 || text.contains("timed out")
1352 || text.contains("implausible response size")
1353}
1354
1355fn encode_single_record(key: Option<Vec<u8>>, value: Option<Vec<u8>>) -> Result<Bytes> {
1357 let timestamp = crate::utils::clock::SystemTime::now()
1358 .duration_since(crate::utils::clock::UNIX_EPOCH)
1359 .map(|d| d.as_millis() as i64)
1360 .unwrap_or(-1);
1361
1362 let records = vec![Record {
1363 transactional: false,
1364 control: false,
1365 partition_leader_epoch: 0,
1366 producer_id: -1,
1367 producer_epoch: -1,
1368 timestamp_type: TimestampType::Creation,
1369 offset: 0,
1370 sequence: 0,
1371 timestamp,
1372 key: key.map(Bytes::from),
1373 value: value.map(Bytes::from),
1374 headers: Default::default(),
1375 }];
1376
1377 let mut buf = Vec::new();
1378 RecordBatchEncoder::encode_with_custom_compression::<_, _, Compressor>(
1379 &mut buf,
1380 &records,
1381 &RecordEncodeOptions {
1382 version: 2,
1383 compression: Compression::None,
1384 },
1385 None,
1386 )
1387 .context("encoding the record batch")?;
1388 Ok(Bytes::from(buf))
1389}
1390
1391fn record_to_json(record: &Record) -> serde_json::Value {
1396 let (key, key_encoding) = match record.key.as_ref() {
1397 Some(k) => {
1398 let (v, e) = encode_field(k.as_ref(), MAX_EVENT_VALUE_BYTES);
1399 (v, Some(e))
1400 }
1401 None => (serde_json::Value::Null, None),
1402 };
1403 let (value, value_encoding) = match record.value.as_ref() {
1404 Some(v) => {
1405 let (val, e) = encode_field(v.as_ref(), MAX_EVENT_VALUE_BYTES);
1406 (val, Some(e))
1407 }
1408 None => (serde_json::Value::Null, None),
1409 };
1410
1411 serde_json::json!({
1412 "offset": record.offset,
1413 "timestamp": record.timestamp,
1414 "key": key,
1415 "key_encoding": key_encoding,
1416 "value": value,
1417 "value_encoding": value_encoding,
1418 })
1419}
1420
1421fn metadata_to_json(response: &MetadataResponse) -> serde_json::Value {
1423 let brokers: Vec<serde_json::Value> = response
1424 .brokers
1425 .iter()
1426 .map(|b| {
1427 serde_json::json!({
1428 "node_id": b.node_id.0,
1429 "host": b.host.to_string(),
1430 "port": b.port,
1431 })
1432 })
1433 .collect();
1434
1435 let topics: Vec<serde_json::Value> = response
1436 .topics
1437 .iter()
1438 .map(|t| {
1439 serde_json::json!({
1440 "name": t.name.as_ref().map(|n| n.to_string()),
1441 "error_code": t.error_code,
1442 "error_name": kafka_error_name(t.error_code),
1443 "partitions": t.partitions.iter().map(|p| serde_json::json!({
1444 "partition": p.partition_index,
1445 "leader": p.leader_id.0,
1446 "error_code": p.error_code,
1447 "replicas": p.replica_nodes.iter().map(|r| r.0).collect::<Vec<_>>(),
1448 })).collect::<Vec<_>>(),
1449 })
1450 })
1451 .collect();
1452
1453 serde_json::json!({
1454 "cluster_id": response.cluster_id.as_ref().map(|c| c.to_string()),
1455 "controller_id": response.controller_id.0,
1456 "brokers": brokers,
1457 "topics": topics,
1458 })
1459}
1460
1461fn str_field(data: &serde_json::Value, key: &str) -> Result<String> {
1462 data.get(key)
1463 .and_then(|v| v.as_str())
1464 .map(|s| s.to_string())
1465 .with_context(|| format!("action data is missing '{key}'"))
1466}
1467
1468fn i64_field(data: &serde_json::Value, key: &str) -> Result<i64> {
1469 data.get(key)
1470 .and_then(|v| v.as_i64())
1471 .with_context(|| format!("action data is missing '{key}'"))
1472}