llm_optimizer_processor/kafka/
sink.rs

1//! Kafka Sink for publishing processed results
2//!
3//! This module provides a high-performance, production-ready Kafka producer
4//! with support for batching, compression, transactions, and exactly-once semantics.
5
6use crate::error::{ProcessorError, Result};
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use rdkafka::config::ClientConfig;
10use rdkafka::error::KafkaError;
11use rdkafka::message::{Header, OwnedHeaders};
12use rdkafka::producer::{FutureProducer, FutureRecord, Producer};
13use rdkafka::util::Timeout;
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
17use std::sync::Arc;
18use std::time::{Duration, Instant};
19use tokio::sync::{Mutex, RwLock, Semaphore};
20use tokio::time::sleep;
21use tracing::{debug, error, info, warn};
22
23/// Maximum number of messages to batch together
24const DEFAULT_BATCH_SIZE: usize = 100;
25
26/// Default timeout for sending messages
27const DEFAULT_SEND_TIMEOUT_MS: u64 = 30000;
28
29/// Default number of retries for failed sends
30const DEFAULT_MAX_RETRIES: u32 = 5;
31
32/// Default base delay for exponential backoff (milliseconds)
33const DEFAULT_BASE_BACKOFF_MS: u64 = 100;
34
35/// Configuration for Kafka sink
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct KafkaSinkConfig {
38    /// Kafka bootstrap servers
39    pub brokers: String,
40
41    /// Default topic for messages
42    pub topic: String,
43
44    /// Client ID for this producer
45    pub client_id: String,
46
47    /// Batch size for batching messages
48    #[serde(default = "default_batch_size")]
49    pub batch_size: usize,
50
51    /// Timeout for sending messages (milliseconds)
52    #[serde(default = "default_send_timeout")]
53    pub send_timeout_ms: u64,
54
55    /// Enable idempotent producer
56    #[serde(default = "default_true")]
57    pub enable_idempotence: bool,
58
59    /// Enable transactions for exactly-once semantics
60    #[serde(default)]
61    pub enable_transactions: bool,
62
63    /// Transaction ID (required if transactions enabled)
64    pub transactional_id: Option<String>,
65
66    /// Compression type (none, gzip, snappy, lz4, zstd)
67    #[serde(default = "default_compression")]
68    pub compression_type: String,
69
70    /// Acknowledgment level (0, 1, all)
71    #[serde(default = "default_acks")]
72    pub acks: String,
73
74    /// Maximum retries for failed sends
75    #[serde(default = "default_max_retries")]
76    pub max_retries: u32,
77
78    /// Base delay for exponential backoff (milliseconds)
79    #[serde(default = "default_base_backoff")]
80    pub base_backoff_ms: u64,
81
82    /// Maximum in-flight requests
83    #[serde(default = "default_max_in_flight")]
84    pub max_in_flight_requests: usize,
85
86    /// Linger time for batching (milliseconds)
87    #[serde(default)]
88    pub linger_ms: u64,
89
90    /// Buffer memory size (bytes)
91    #[serde(default = "default_buffer_memory")]
92    pub buffer_memory: usize,
93
94    /// Enable circuit breaker
95    #[serde(default = "default_true")]
96    pub enable_circuit_breaker: bool,
97
98    /// Circuit breaker failure threshold
99    #[serde(default = "default_circuit_breaker_threshold")]
100    pub circuit_breaker_threshold: u32,
101
102    /// Circuit breaker timeout (seconds)
103    #[serde(default = "default_circuit_breaker_timeout")]
104    pub circuit_breaker_timeout_secs: u64,
105}
106
107fn default_batch_size() -> usize { DEFAULT_BATCH_SIZE }
108fn default_send_timeout() -> u64 { DEFAULT_SEND_TIMEOUT_MS }
109fn default_true() -> bool { true }
110fn default_compression() -> String { "snappy".to_string() }
111fn default_acks() -> String { "all".to_string() }
112fn default_max_retries() -> u32 { DEFAULT_MAX_RETRIES }
113fn default_base_backoff() -> u64 { DEFAULT_BASE_BACKOFF_MS }
114fn default_max_in_flight() -> usize { 5 }
115fn default_buffer_memory() -> usize { 33554432 } // 32MB
116fn default_circuit_breaker_threshold() -> u32 { 5 }
117fn default_circuit_breaker_timeout() -> u64 { 60 }
118
119impl Default for KafkaSinkConfig {
120    fn default() -> Self {
121        Self {
122            brokers: "localhost:9092".to_string(),
123            topic: "llm-optimizer-results".to_string(),
124            client_id: "llm-optimizer-producer".to_string(),
125            batch_size: DEFAULT_BATCH_SIZE,
126            send_timeout_ms: DEFAULT_SEND_TIMEOUT_MS,
127            enable_idempotence: true,
128            enable_transactions: false,
129            transactional_id: None,
130            compression_type: "snappy".to_string(),
131            acks: "all".to_string(),
132            max_retries: DEFAULT_MAX_RETRIES,
133            base_backoff_ms: DEFAULT_BASE_BACKOFF_MS,
134            max_in_flight_requests: 5,
135            linger_ms: 10,
136            buffer_memory: 33554432,
137            enable_circuit_breaker: true,
138            circuit_breaker_threshold: 5,
139            circuit_breaker_timeout_secs: 60,
140        }
141    }
142}
143
144/// Partitioning strategy for messages
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
146pub enum PartitionStrategy {
147    /// Use key hash for partitioning
148    KeyHash,
149    /// Round-robin partitioning
150    RoundRobin,
151    /// Custom partition specified per message
152    Custom,
153    /// All messages to partition 0
154    Single,
155}
156
157/// Delivery guarantee level
158#[derive(Debug, Clone, Copy, PartialEq, Eq)]
159pub enum DeliveryGuarantee {
160    /// At-most-once (fire and forget)
161    AtMostOnce,
162    /// At-least-once (retry on failure)
163    AtLeastOnce,
164    /// Exactly-once (transactions required)
165    ExactlyOnce,
166}
167
168/// Message to be sent to Kafka
169#[derive(Debug, Clone)]
170pub struct KafkaMessage<T> {
171    /// Message payload
172    pub payload: T,
173    /// Message key (optional)
174    pub key: Option<String>,
175    /// Target topic (uses default if None)
176    pub topic: Option<String>,
177    /// Target partition (for custom partitioning)
178    pub partition: Option<i32>,
179    /// Message headers
180    pub headers: HashMap<String, String>,
181    /// Timestamp
182    pub timestamp: Option<DateTime<Utc>>,
183}
184
185impl<T> KafkaMessage<T> {
186    /// Create a new message with payload
187    pub fn new(payload: T) -> Self {
188        Self {
189            payload,
190            key: None,
191            topic: None,
192            partition: None,
193            headers: HashMap::new(),
194            timestamp: None,
195        }
196    }
197
198    /// Set message key
199    pub fn with_key(mut self, key: String) -> Self {
200        self.key = Some(key);
201        self
202    }
203
204    /// Set topic
205    pub fn with_topic(mut self, topic: String) -> Self {
206        self.topic = Some(topic);
207        self
208    }
209
210    /// Set partition
211    pub fn with_partition(mut self, partition: i32) -> Self {
212        self.partition = Some(partition);
213        self
214    }
215
216    /// Add header
217    pub fn with_header(mut self, key: String, value: String) -> Self {
218        self.headers.insert(key, value);
219        self
220    }
221
222    /// Set timestamp
223    pub fn with_timestamp(mut self, timestamp: DateTime<Utc>) -> Self {
224        self.timestamp = Some(timestamp);
225        self
226    }
227}
228
229/// Production metrics
230#[derive(Debug, Default, Clone)]
231pub struct SinkMetrics {
232    /// Total messages sent successfully
233    pub messages_sent: u64,
234    /// Total messages failed
235    pub messages_failed: u64,
236    /// Total bytes sent
237    pub bytes_sent: u64,
238    /// Total send attempts
239    pub send_attempts: u64,
240    /// Total retries
241    pub retries: u64,
242    /// Average send latency (microseconds)
243    pub avg_latency_us: u64,
244    /// Maximum send latency (microseconds)
245    pub max_latency_us: u64,
246    /// Circuit breaker trips
247    pub circuit_breaker_trips: u64,
248    /// Last error timestamp
249    pub last_error: Option<DateTime<Utc>>,
250    /// Last error message
251    pub last_error_msg: Option<String>,
252}
253
254/// Internal metrics tracking
255struct MetricsTracker {
256    messages_sent: AtomicU64,
257    messages_failed: AtomicU64,
258    bytes_sent: AtomicU64,
259    send_attempts: AtomicU64,
260    retries: AtomicU64,
261    total_latency_us: AtomicU64,
262    max_latency_us: AtomicU64,
263    circuit_breaker_trips: AtomicU64,
264    last_error: Mutex<Option<(DateTime<Utc>, String)>>,
265}
266
267impl MetricsTracker {
268    fn new() -> Self {
269        Self {
270            messages_sent: AtomicU64::new(0),
271            messages_failed: AtomicU64::new(0),
272            bytes_sent: AtomicU64::new(0),
273            send_attempts: AtomicU64::new(0),
274            retries: AtomicU64::new(0),
275            total_latency_us: AtomicU64::new(0),
276            max_latency_us: AtomicU64::new(0),
277            circuit_breaker_trips: AtomicU64::new(0),
278            last_error: Mutex::new(None),
279        }
280    }
281
282    async fn record_success(&self, latency_us: u64, bytes: u64) {
283        self.messages_sent.fetch_add(1, Ordering::Relaxed);
284        self.bytes_sent.fetch_add(bytes, Ordering::Relaxed);
285        self.total_latency_us.fetch_add(latency_us, Ordering::Relaxed);
286
287        let mut max = self.max_latency_us.load(Ordering::Relaxed);
288        while latency_us > max {
289            match self.max_latency_us.compare_exchange_weak(
290                max,
291                latency_us,
292                Ordering::Relaxed,
293                Ordering::Relaxed,
294            ) {
295                Ok(_) => break,
296                Err(x) => max = x,
297            }
298        }
299    }
300
301    async fn record_failure(&self, error: String) {
302        self.messages_failed.fetch_add(1, Ordering::Relaxed);
303        let mut last_error = self.last_error.lock().await;
304        *last_error = Some((Utc::now(), error));
305    }
306
307    fn record_attempt(&self) {
308        self.send_attempts.fetch_add(1, Ordering::Relaxed);
309    }
310
311    fn record_retry(&self) {
312        self.retries.fetch_add(1, Ordering::Relaxed);
313    }
314
315    fn record_circuit_breaker_trip(&self) {
316        self.circuit_breaker_trips.fetch_add(1, Ordering::Relaxed);
317    }
318
319    async fn snapshot(&self) -> SinkMetrics {
320        let messages_sent = self.messages_sent.load(Ordering::Relaxed);
321        let total_latency = self.total_latency_us.load(Ordering::Relaxed);
322        let avg_latency = if messages_sent > 0 {
323            total_latency / messages_sent
324        } else {
325            0
326        };
327
328        let last_error = self.last_error.lock().await;
329        let (last_error_time, last_error_msg) = last_error.as_ref()
330            .map(|(t, m)| (Some(*t), Some(m.clone())))
331            .unwrap_or((None, None));
332
333        SinkMetrics {
334            messages_sent,
335            messages_failed: self.messages_failed.load(Ordering::Relaxed),
336            bytes_sent: self.bytes_sent.load(Ordering::Relaxed),
337            send_attempts: self.send_attempts.load(Ordering::Relaxed),
338            retries: self.retries.load(Ordering::Relaxed),
339            avg_latency_us: avg_latency,
340            max_latency_us: self.max_latency_us.load(Ordering::Relaxed),
341            circuit_breaker_trips: self.circuit_breaker_trips.load(Ordering::Relaxed),
342            last_error: last_error_time,
343            last_error_msg,
344        }
345    }
346}
347
348/// Circuit breaker state
349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
350enum CircuitState {
351    Closed,
352    Open,
353    HalfOpen,
354}
355
356/// Circuit breaker for fault tolerance
357struct CircuitBreaker {
358    state: RwLock<CircuitState>,
359    failure_count: AtomicU64,
360    threshold: u32,
361    timeout: Duration,
362    last_failure_time: RwLock<Option<Instant>>,
363}
364
365impl CircuitBreaker {
366    fn new(threshold: u32, timeout: Duration) -> Self {
367        Self {
368            state: RwLock::new(CircuitState::Closed),
369            failure_count: AtomicU64::new(0),
370            threshold,
371            timeout,
372            last_failure_time: RwLock::new(None),
373        }
374    }
375
376    async fn can_attempt(&self) -> bool {
377        let state = *self.state.read().await;
378
379        match state {
380            CircuitState::Closed => true,
381            CircuitState::HalfOpen => true,
382            CircuitState::Open => {
383                // Check if timeout has elapsed
384                if let Some(last_failure) = *self.last_failure_time.read().await {
385                    if last_failure.elapsed() >= self.timeout {
386                        // Transition to half-open
387                        *self.state.write().await = CircuitState::HalfOpen;
388                        self.failure_count.store(0, Ordering::Relaxed);
389                        true
390                    } else {
391                        false
392                    }
393                } else {
394                    false
395                }
396            }
397        }
398    }
399
400    async fn record_success(&self) {
401        let mut state = self.state.write().await;
402        *state = CircuitState::Closed;
403        self.failure_count.store(0, Ordering::Relaxed);
404    }
405
406    async fn record_failure(&self) {
407        let failures = self.failure_count.fetch_add(1, Ordering::Relaxed) + 1;
408        *self.last_failure_time.write().await = Some(Instant::now());
409
410        if failures >= self.threshold as u64 {
411            let mut state = self.state.write().await;
412            *state = CircuitState::Open;
413        }
414    }
415
416    async fn is_open(&self) -> bool {
417        *self.state.read().await == CircuitState::Open
418    }
419}
420
421/// Serializer trait for converting messages to bytes
422#[async_trait]
423pub trait MessageSerializer<T>: Send + Sync {
424    /// Serialize message to bytes
425    async fn serialize(&self, message: &T) -> Result<Vec<u8>>;
426
427    /// Get content type identifier
428    fn content_type(&self) -> &str;
429}
430
431/// JSON serializer
432pub struct JsonSerializer;
433
434#[async_trait]
435impl<T: Serialize + Send + Sync> MessageSerializer<T> for JsonSerializer {
436    async fn serialize(&self, message: &T) -> Result<Vec<u8>> {
437        serde_json::to_vec(message).map_err(|e| ProcessorError::Serialization(e.to_string()))
438    }
439
440    fn content_type(&self) -> &str {
441        "application/json"
442    }
443}
444
445/// Bincode serializer
446pub struct BincodeSerializer;
447
448#[async_trait]
449impl<T: Serialize + Send + Sync> MessageSerializer<T> for BincodeSerializer {
450    async fn serialize(&self, message: &T) -> Result<Vec<u8>> {
451        bincode::serialize(message).map_err(|e| ProcessorError::Serialization(e.to_string()))
452    }
453
454    fn content_type(&self) -> &str {
455        "application/octet-stream"
456    }
457}
458
459/// Kafka sink for publishing messages
460pub struct KafkaSink<T> {
461    /// Kafka producer
462    producer: Arc<FutureProducer>,
463
464    /// Configuration
465    config: KafkaSinkConfig,
466
467    /// Message serializer
468    serializer: Arc<dyn MessageSerializer<T>>,
469
470    /// Partitioning strategy
471    partition_strategy: PartitionStrategy,
472
473    /// Delivery guarantee
474    delivery_guarantee: DeliveryGuarantee,
475
476    /// Metrics tracker
477    metrics: Arc<MetricsTracker>,
478
479    /// Circuit breaker
480    circuit_breaker: Option<Arc<CircuitBreaker>>,
481
482    /// Semaphore for controlling in-flight requests
483    in_flight_semaphore: Arc<Semaphore>,
484
485    /// Whether transactions are active
486    in_transaction: AtomicBool,
487
488    /// Shutdown flag
489    shutdown: AtomicBool,
490}
491
492impl<T> KafkaSink<T>
493where
494    T: Serialize + Send + Sync + 'static,
495{
496    /// Create a new Kafka sink with JSON serialization
497    pub async fn new(config: KafkaSinkConfig) -> Result<Self> {
498        Self::with_serializer(config, Arc::new(JsonSerializer)).await
499    }
500
501    /// Create a new Kafka sink with custom serializer
502    pub async fn with_serializer(
503        config: KafkaSinkConfig,
504        serializer: Arc<dyn MessageSerializer<T>>,
505    ) -> Result<Self> {
506        // Validate configuration
507        if config.enable_transactions && config.transactional_id.is_none() {
508            return Err(ProcessorError::Configuration {
509                source: "transactional_id required when transactions are enabled".into(),
510            });
511        }
512
513        // Build Kafka producer config
514        let mut client_config = ClientConfig::new();
515        client_config
516            .set("bootstrap.servers", &config.brokers)
517            .set("client.id", &config.client_id)
518            .set("compression.type", &config.compression_type)
519            .set("acks", &config.acks)
520            .set("linger.ms", config.linger_ms.to_string())
521            .set("batch.size", config.batch_size.to_string())
522            .set("buffer.memory", config.buffer_memory.to_string())
523            .set("max.in.flight.requests.per.connection", config.max_in_flight_requests.to_string())
524            .set("retries", config.max_retries.to_string());
525
526        if config.enable_idempotence {
527            client_config.set("enable.idempotence", "true");
528        }
529
530        if config.enable_transactions {
531            if let Some(ref txn_id) = config.transactional_id {
532                client_config.set("transactional.id", txn_id);
533            }
534        }
535
536        // Create producer
537        let producer: FutureProducer = client_config
538            .create()
539            .map_err(|e| ProcessorError::Configuration {
540                source: Box::new(e),
541            })?;
542
543        // Initialize transactions if enabled
544        if config.enable_transactions {
545            producer.init_transactions(Timeout::After(Duration::from_secs(30)))
546                .map_err(|e| ProcessorError::Configuration {
547                    source: Box::new(e),
548                })?;
549        }
550
551        let circuit_breaker = if config.enable_circuit_breaker {
552            Some(Arc::new(CircuitBreaker::new(
553                config.circuit_breaker_threshold,
554                Duration::from_secs(config.circuit_breaker_timeout_secs),
555            )))
556        } else {
557            None
558        };
559
560        Ok(Self {
561            producer: Arc::new(producer),
562            config,
563            serializer,
564            partition_strategy: PartitionStrategy::KeyHash,
565            delivery_guarantee: DeliveryGuarantee::AtLeastOnce,
566            metrics: Arc::new(MetricsTracker::new()),
567            circuit_breaker,
568            in_flight_semaphore: Arc::new(Semaphore::new(5)),
569            in_transaction: AtomicBool::new(false),
570            shutdown: AtomicBool::new(false),
571        })
572    }
573
574    /// Set partitioning strategy
575    pub fn with_partition_strategy(mut self, strategy: PartitionStrategy) -> Self {
576        self.partition_strategy = strategy;
577        self
578    }
579
580    /// Set delivery guarantee
581    pub fn with_delivery_guarantee(mut self, guarantee: DeliveryGuarantee) -> Self {
582        self.delivery_guarantee = guarantee;
583        self
584    }
585
586    /// Send a single message
587    pub async fn send(&self, message: KafkaMessage<T>) -> Result<()> {
588        if self.shutdown.load(Ordering::Relaxed) {
589            return Err(ProcessorError::Unexpected("sink is shut down".to_string()));
590        }
591
592        // Check circuit breaker
593        if let Some(ref cb) = self.circuit_breaker {
594            if !cb.can_attempt().await {
595                self.metrics.record_circuit_breaker_trip();
596                return Err(ProcessorError::Unexpected(
597                    "circuit breaker is open".to_string()
598                ));
599            }
600        }
601
602        // Serialize message
603        let payload_bytes = self.serializer.serialize(&message.payload).await?;
604        let payload_size = payload_bytes.len() as u64;
605
606        // Send with retry logic
607        let result = self.send_with_retry(message, payload_bytes).await;
608
609        // Update circuit breaker
610        if let Some(ref cb) = self.circuit_breaker {
611            match result {
612                Ok(_) => cb.record_success().await,
613                Err(_) => cb.record_failure().await,
614            }
615        }
616
617        // Update metrics
618        match result {
619            Ok(latency_us) => {
620                self.metrics.record_success(latency_us, payload_size).await;
621                Ok(())
622            }
623            Err(e) => {
624                self.metrics.record_failure(e.to_string()).await;
625                Err(e)
626            }
627        }
628    }
629
630    /// Send a batch of messages
631    pub async fn send_batch(&self, messages: Vec<KafkaMessage<T>>) -> Result<Vec<Result<()>>> {
632        if self.shutdown.load(Ordering::Relaxed) {
633            return Err(ProcessorError::Unexpected("sink is shut down".to_string()));
634        }
635
636        let mut results = Vec::with_capacity(messages.len());
637
638        // Process messages in parallel with concurrency control
639        let mut handles = Vec::new();
640
641        for message in messages {
642            let permit = self.in_flight_semaphore.clone().acquire_owned().await
643                .map_err(|e| ProcessorError::Unexpected(e.to_string()))?;
644
645            let sink = self.clone_for_send();
646            let handle = tokio::spawn(async move {
647                let result = sink.send(message).await;
648                drop(permit);
649                result
650            });
651
652            handles.push(handle);
653        }
654
655        // Collect results
656        for handle in handles {
657            let result = handle.await
658                .map_err(|e| ProcessorError::Unexpected(e.to_string()))?;
659            results.push(result);
660        }
661
662        Ok(results)
663    }
664
665    /// Begin a transaction
666    pub async fn begin_transaction(&self) -> Result<()> {
667        if !self.config.enable_transactions {
668            return Err(ProcessorError::Configuration {
669                source: "transactions not enabled".into(),
670            });
671        }
672
673        if self.in_transaction.swap(true, Ordering::SeqCst) {
674            return Err(ProcessorError::Unexpected(
675                "transaction already in progress".to_string()
676            ));
677        }
678
679        self.producer
680            .begin_transaction()
681            .map_err(|e| ProcessorError::Execution {
682                source: Box::new(e),
683            })?;
684
685        debug!("Transaction begun");
686        Ok(())
687    }
688
689    /// Commit current transaction
690    pub async fn commit_transaction(&self) -> Result<()> {
691        if !self.in_transaction.load(Ordering::SeqCst) {
692            return Err(ProcessorError::Unexpected(
693                "no transaction in progress".to_string()
694            ));
695        }
696
697        self.producer
698            .commit_transaction(Timeout::After(Duration::from_secs(30)))
699            .map_err(|e| ProcessorError::Execution {
700                source: Box::new(e),
701            })?;
702
703        self.in_transaction.store(false, Ordering::SeqCst);
704        debug!("Transaction committed");
705        Ok(())
706    }
707
708    /// Abort current transaction
709    pub async fn abort_transaction(&self) -> Result<()> {
710        if !self.in_transaction.load(Ordering::SeqCst) {
711            return Err(ProcessorError::Unexpected(
712                "no transaction in progress".to_string()
713            ));
714        }
715
716        self.producer
717            .abort_transaction(Timeout::After(Duration::from_secs(30)))
718            .map_err(|e| ProcessorError::Execution {
719                source: Box::new(e),
720            })?;
721
722        self.in_transaction.store(false, Ordering::SeqCst);
723        warn!("Transaction aborted");
724        Ok(())
725    }
726
727    /// Flush pending messages
728    pub async fn flush(&self) -> Result<()> {
729        info!("Flushing pending messages...");
730
731        self.producer
732            .flush(Timeout::After(Duration::from_millis(self.config.send_timeout_ms)))
733            .map_err(|e| ProcessorError::Execution {
734                source: Box::new(e),
735            })?;
736
737        info!("Flush completed");
738        Ok(())
739    }
740
741    /// Get current metrics
742    pub async fn metrics(&self) -> SinkMetrics {
743        self.metrics.snapshot().await
744    }
745
746    /// Graceful shutdown
747    pub async fn close(&self) -> Result<()> {
748        info!("Shutting down Kafka sink...");
749
750        self.shutdown.store(true, Ordering::Relaxed);
751
752        // Abort any active transaction
753        if self.in_transaction.load(Ordering::SeqCst) {
754            if let Err(e) = self.abort_transaction().await {
755                error!("Failed to abort transaction during shutdown: {}", e);
756            }
757        }
758
759        // Flush pending messages
760        if let Err(e) = self.flush().await {
761            error!("Failed to flush during shutdown: {}", e);
762        }
763
764        let metrics = self.metrics().await;
765        info!(
766            "Kafka sink shut down. Sent: {}, Failed: {}, Avg latency: {}us",
767            metrics.messages_sent,
768            metrics.messages_failed,
769            metrics.avg_latency_us
770        );
771
772        Ok(())
773    }
774
775    /// Internal: Send message with retry logic
776    async fn send_with_retry(
777        &self,
778        message: KafkaMessage<T>,
779        payload_bytes: Vec<u8>,
780    ) -> Result<u64> {
781        let mut retries = 0;
782        let max_retries = self.config.max_retries;
783
784        loop {
785            self.metrics.record_attempt();
786            let start = Instant::now();
787
788            let topic = message.topic.as_ref()
789                .unwrap_or(&self.config.topic);
790
791            // Build Kafka record
792            let mut record = FutureRecord::to(topic)
793                .payload(&payload_bytes);
794
795            if let Some(ref key) = message.key {
796                record = record.key(key);
797            }
798
799            if let Some(partition) = message.partition {
800                record = record.partition(partition);
801            }
802
803            if let Some(timestamp) = message.timestamp {
804                record = record.timestamp(timestamp.timestamp_millis());
805            }
806
807            // Add headers
808            if !message.headers.is_empty() {
809                let mut headers = OwnedHeaders::new();
810                for (k, v) in &message.headers {
811                    headers = headers.insert(Header {
812                        key: k,
813                        value: Some(v.as_bytes()),
814                    });
815                }
816                // Add content type
817                headers = headers.insert(Header {
818                    key: "content-type",
819                    value: Some(self.serializer.content_type().as_bytes()),
820                });
821                record = record.headers(headers);
822            }
823
824            // Send message
825            let result = self.producer
826                .send(
827                    record,
828                    Timeout::After(Duration::from_millis(self.config.send_timeout_ms)),
829                )
830                .await;
831
832            let latency_us = start.elapsed().as_micros() as u64;
833
834            match result {
835                Ok(_) => {
836                    if retries > 0 {
837                        debug!("Message sent successfully after {} retries", retries);
838                    }
839                    return Ok(latency_us);
840                }
841                Err((err, _)) => {
842                    if !self.should_retry(&err, retries, max_retries) {
843                        error!("Failed to send message after {} retries: {}", retries, err);
844                        return Err(ProcessorError::Execution {
845                            source: Box::new(err),
846                        });
847                    }
848
849                    retries += 1;
850                    self.metrics.record_retry();
851
852                    let backoff = self.calculate_backoff(retries);
853                    warn!(
854                        "Send failed (attempt {}/{}), retrying in {:?}: {}",
855                        retries, max_retries + 1, backoff, err
856                    );
857                    sleep(backoff).await;
858                }
859            }
860        }
861    }
862
863    /// Check if error is retriable
864    fn should_retry(&self, error: &KafkaError, retries: u32, max_retries: u32) -> bool {
865        if retries >= max_retries {
866            return false;
867        }
868
869        // Check if error is retriable
870        match error {
871            KafkaError::MessageProduction(rdkafka::types::RDKafkaErrorCode::QueueFull) => true,
872            KafkaError::MessageProduction(rdkafka::types::RDKafkaErrorCode::NetworkException) => true,
873            KafkaError::MessageProduction(rdkafka::types::RDKafkaErrorCode::RequestTimedOut) => true,
874            KafkaError::MessageProduction(rdkafka::types::RDKafkaErrorCode::NotLeaderForPartition) => true,
875            _ => false,
876        }
877    }
878
879    /// Calculate exponential backoff delay
880    fn calculate_backoff(&self, retry_count: u32) -> Duration {
881        let base_ms = self.config.base_backoff_ms;
882        let backoff_ms = base_ms * 2u64.pow(retry_count.min(10));
883        let max_backoff_ms = 60_000; // 1 minute max
884        Duration::from_millis(backoff_ms.min(max_backoff_ms))
885    }
886
887    /// Clone for parallel sends (shares producer and metrics)
888    fn clone_for_send(&self) -> Self {
889        Self {
890            producer: Arc::clone(&self.producer),
891            config: self.config.clone(),
892            serializer: Arc::clone(&self.serializer),
893            partition_strategy: self.partition_strategy,
894            delivery_guarantee: self.delivery_guarantee,
895            metrics: Arc::clone(&self.metrics),
896            circuit_breaker: self.circuit_breaker.as_ref().map(Arc::clone),
897            in_flight_semaphore: Arc::clone(&self.in_flight_semaphore),
898            in_transaction: AtomicBool::new(false),
899            shutdown: AtomicBool::new(false),
900        }
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907
908    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
909    struct TestMessage {
910        id: String,
911        value: i32,
912    }
913
914    #[tokio::test]
915    async fn test_kafka_message_builder() {
916        let msg = KafkaMessage::new(TestMessage {
917            id: "test".to_string(),
918            value: 42,
919        })
920        .with_key("key1".to_string())
921        .with_topic("test-topic".to_string())
922        .with_header("trace-id".to_string(), "12345".to_string());
923
924        assert_eq!(msg.key, Some("key1".to_string()));
925        assert_eq!(msg.topic, Some("test-topic".to_string()));
926        assert_eq!(msg.headers.get("trace-id"), Some(&"12345".to_string()));
927    }
928
929    #[tokio::test]
930    async fn test_json_serializer() {
931        let serializer = JsonSerializer;
932        let msg = TestMessage {
933            id: "test".to_string(),
934            value: 42,
935        };
936
937        let bytes = MessageSerializer::serialize(&serializer, &msg).await.unwrap();
938        let deserialized: TestMessage = serde_json::from_slice(&bytes).unwrap();
939
940        assert_eq!(msg, deserialized);
941        assert_eq!(MessageSerializer::<TestMessage>::content_type(&serializer), "application/json");
942    }
943
944    #[tokio::test]
945    async fn test_bincode_serializer() {
946        let serializer = BincodeSerializer;
947        let msg = TestMessage {
948            id: "test".to_string(),
949            value: 42,
950        };
951
952        let bytes = MessageSerializer::serialize(&serializer, &msg).await.unwrap();
953        let deserialized: TestMessage = bincode::deserialize(&bytes).unwrap();
954
955        assert_eq!(msg, deserialized);
956        assert_eq!(MessageSerializer::<TestMessage>::content_type(&serializer), "application/octet-stream");
957    }
958
959    #[tokio::test]
960    async fn test_metrics_tracker() {
961        let tracker = MetricsTracker::new();
962
963        tracker.record_success(1000, 100).await;
964        tracker.record_success(2000, 200).await;
965        tracker.record_failure("test error".to_string()).await;
966        tracker.record_attempt();
967        tracker.record_retry();
968
969        let metrics = tracker.snapshot().await;
970
971        assert_eq!(metrics.messages_sent, 2);
972        assert_eq!(metrics.messages_failed, 1);
973        assert_eq!(metrics.bytes_sent, 300);
974        assert_eq!(metrics.send_attempts, 1);
975        assert_eq!(metrics.retries, 1);
976        assert_eq!(metrics.avg_latency_us, 1500);
977        assert_eq!(metrics.max_latency_us, 2000);
978    }
979
980    #[tokio::test]
981    async fn test_circuit_breaker() {
982        let cb = CircuitBreaker::new(3, Duration::from_millis(100));
983
984        // Initially closed
985        assert!(cb.can_attempt().await);
986        assert!(!cb.is_open().await);
987
988        // Record failures
989        cb.record_failure().await;
990        cb.record_failure().await;
991        assert!(!cb.is_open().await);
992
993        cb.record_failure().await;
994        assert!(cb.is_open().await);
995        assert!(!cb.can_attempt().await);
996
997        // Wait for timeout
998        tokio::time::sleep(Duration::from_millis(150)).await;
999
1000        // Should transition to half-open
1001        assert!(cb.can_attempt().await);
1002
1003        // Success closes circuit
1004        cb.record_success().await;
1005        assert!(!cb.is_open().await);
1006    }
1007
1008    #[test]
1009    fn test_config_defaults() {
1010        let config = KafkaSinkConfig::default();
1011
1012        assert_eq!(config.brokers, "localhost:9092");
1013        assert_eq!(config.batch_size, DEFAULT_BATCH_SIZE);
1014        assert!(config.enable_idempotence);
1015        assert!(!config.enable_transactions);
1016        assert_eq!(config.compression_type, "snappy");
1017        assert_eq!(config.acks, "all");
1018    }
1019
1020    #[test]
1021    fn test_partition_strategy() {
1022        assert_eq!(PartitionStrategy::KeyHash, PartitionStrategy::KeyHash);
1023        assert_ne!(PartitionStrategy::KeyHash, PartitionStrategy::RoundRobin);
1024    }
1025
1026    #[test]
1027    fn test_delivery_guarantee() {
1028        assert_eq!(DeliveryGuarantee::AtLeastOnce, DeliveryGuarantee::AtLeastOnce);
1029        assert_ne!(DeliveryGuarantee::AtLeastOnce, DeliveryGuarantee::ExactlyOnce);
1030    }
1031}
llm_optimizer_processor/kafka/sink.rs

llm_optimizer_processor/kafka/
sink.rs
