support aio kafka metrics

Author: idanasulin2706

superclient/agent/interceptor.py

@@ -3,6 +3,7 @@
 import os
 import uuid
 from typing import Any, Dict
+import importlib
 
 from ..util.logger import get_logger
 from ..util.config import get_topics_list, is_disabled
@@ -245,8 +246,102 @@ async def stop_patch(*a, **kw):
                     await original_stop(*a, **kw)
 
                 self.stop = stop_patch
+                
+                # Patch the SendProduceReqHandler.create_request method for this specific producer
+                sender_mod = importlib.import_module("aiokafka.producer.sender")
+                
+                # Only patch once globally
+                if not hasattr(sender_mod.SendProduceReqHandler, '_superstream_patched'):
+                    orig_create_request = sender_mod.SendProduceReqHandler.create_request
+                    
+                    def create_request_with_metrics(self_handler):
+                        # Call the original method to get the request, but collect metrics
+                        # self_handler._batches: Dict[TopicPartition, MessageBatch]
+                        
+                        # Quick check: if sender has no tracker, it's an internal producer - skip metrics
+                        if not hasattr(self_handler._sender, '_superstream_tracker'):
+                            return orig_create_request(self_handler)
+                            
+                        tracker = self_handler._sender._superstream_tracker
+                        
+                        # Additional check: skip internal producers by client_id
+                        if tracker is None or tracker.client_id.startswith(_SUPERLIB_PREFIX):
+                            return orig_create_request(self_handler)
+                            
+                        # Per-producer totals
+                        total_uncompressed = 0
+                        total_compressed = 0
+                        total_records = 0
+                        topic_stats = {}
+                        for tp, batch in self_handler._batches.items():
+                            # Get record count from the batch
+                            record_count = batch.record_count
+                            
+                            # Get compressed size from the batch buffer
+                            compressed = 0
+                            try:
+                                compressed = len(batch.get_data_buffer())
+                            except Exception:
+                                pass
+                            
+                            # Estimate uncompressed size based on record count
+                            # Since we can't easily access the original message data at this point,
+                            # we'll use a reasonable estimate based on the batch size and record count
+                            if record_count > 0:
+                                # Estimate uncompressed size based on compressed size and typical compression ratios
+                                # This is an approximation since we can't access the original message data
+                                estimated_compression_ratio = 0.7  # Assume 30% compression
+                                uncompressed = int(compressed / estimated_compression_ratio)
+                            else:
+                                uncompressed = 0
+                            
+                            total_uncompressed += uncompressed
+                            total_compressed += compressed
+                            total_records += record_count
+                            # Per-topic
+                            if tp.topic not in topic_stats:
+                                topic_stats[tp.topic] = {'uncompressed': 0, 'compressed': 0, 'records': 0}
+                            topic_stats[tp.topic]['uncompressed'] += uncompressed
+                            topic_stats[tp.topic]['compressed'] += compressed
+                            topic_stats[tp.topic]['records'] += record_count
+                        # Update tracker
+                        if total_records > 0:
+                            tracker._superstream_metrics = getattr(tracker, '_superstream_metrics', {})
+                            m = tracker._superstream_metrics
+                            # Accumulate totals (aggregative counters)
+                            m['outgoing-byte-total'] = m.get('outgoing-byte-total', 0) + total_compressed
+                            m['record-send-total'] = m.get('record-send-total', 0) + total_records
+                            m['uncompressed-byte-total'] = m.get('uncompressed-byte-total', 0) + total_uncompressed
+                            
+                            # Calculate rates from aggregated totals
+                            m['compression-rate-avg'] = (m['outgoing-byte-total'] / m['uncompressed-byte-total']) if m['uncompressed-byte-total'] else 1.0
+                            m['record-size-avg'] = (m['uncompressed-byte-total'] / m['record-send-total']) if m['record-send-total'] else 0
+                            
+                            # Per-topic
+                            m['topics'] = m.get('topics', {})
+                            for topic, stats in topic_stats.items():
+                                t = m['topics'].setdefault(topic, {'byte-total': 0, 'record-send-total': 0, 'uncompressed-total': 0})
+                                # Accumulate totals (aggregative counters)
+                                t['byte-total'] = t.get('byte-total', 0) + stats['compressed']
+                                t['record-send-total'] = t.get('record-send-total', 0) + stats['records']
+                                t['uncompressed-total'] = t.get('uncompressed-total', 0) + stats['uncompressed']
+                                
+                                # Calculate compression rate from aggregated totals
+                                t['compression-rate'] = (t['byte-total'] / t['uncompressed-total']) if t['uncompressed-total'] else 1.0
+                                
+                            tracker._superstream_metrics = m
+                        return orig_create_request(self_handler)
+                    
+                    sender_mod.SendProduceReqHandler.create_request = create_request_with_metrics
+                    sender_mod.SendProduceReqHandler._superstream_patched = True
+                
                 self._superstream_patch = True
             orig_init(self, *args, **kwargs)
+            
+                        # Store tracker reference in the sender for metrics collection
+            if hasattr(self, '_sender'):
+                self._sender._superstream_tracker = tr
+            
             send_clients_msg(tr, error_msg)
 
             # Log success message based on whether defaults were used

superclient/agent/tracker.py

@@ -2,7 +2,7 @@
 
 import threading
 import time
-import uuid
+import uuid as uuid_lib
 from typing import Any, Dict, Optional
 
 from ..util.logger import get_logger
@@ -26,7 +26,7 @@ def __init__(
         topics_env: Optional[list[str]] = None,
         uuid: Optional[str] = None,
     ) -> None:
-        self.uuid = uuid if uuid else str(uuid.uuid4())
+        self.uuid = uuid if uuid else str(uuid_lib.uuid4())
         self.library = lib
         self.producer = producer
         self.bootstrap = bootstrap

superclient/util/metrics.py

@@ -92,11 +92,20 @@ def collect_confluent_producer_metrics(tracker_uuid: str) -> Dict[str, Any]:
         return {}
 
 
-def collect_aiokafka_metrics(producer: Any) -> Dict[str, Any]:
+def collect_aiokafka_producer_metrics(producer: Any) -> Dict[str, Any]:
     """Collect metrics from aiokafka producer."""
     try:
-        # aiokafka producers don't have a metrics() method
-        # Return empty dict as aiokafka doesn't provide metrics
+        # Get the tracker from the sender (where we store it in aiokafka)
+        if hasattr(producer, '_sender') and hasattr(producer._sender, '_superstream_tracker'):
+            tracker = producer._sender._superstream_tracker
+            if hasattr(tracker, '_superstream_metrics'):
+                metrics = tracker._superstream_metrics
+                # Return only the producer-level metrics (not topics)
+                producer_metrics = {}
+                for key in ['outgoing-byte-total', 'compression-rate-avg', 'record-send-total', 'record-size-avg']:
+                    if key in metrics:
+                        producer_metrics[key] = sanitize_metric_value(metrics[key])
+                return producer_metrics
         return {}
     except Exception as e:
         logger.error("[ERR-308] Failed to collect aiokafka producer metrics: {}", e)
@@ -151,8 +160,18 @@ def collect_confluent_topic_metrics(tracker_uuid: str) -> Dict[str, Any]:
 def collect_aiokafka_topic_metrics(producer: Any) -> Dict[str, Any]:
     """Collect topic metrics from aiokafka producer."""
     try:
-        # aiokafka producers don't have a metrics() method
-        # Return empty dict as aiokafka doesn't provide topic metrics
+        # Get the tracker from the sender (where we store it in aiokafka)
+        if hasattr(producer, '_sender') and hasattr(producer._sender, '_superstream_tracker'):
+            tracker = producer._sender._superstream_tracker
+            if hasattr(tracker, '_superstream_metrics') and 'topics' in tracker._superstream_metrics:
+                topic_metrics = {}
+                for topic_name, topic_data in tracker._superstream_metrics['topics'].items():
+                    # Return only the required topic metrics
+                    topic_metrics[topic_name] = {}
+                    for key in ['byte-total', 'compression-rate', 'record-send-total']:
+                        if key in topic_data:
+                            topic_metrics[topic_name][key] = sanitize_metric_value(topic_data[key])
+                return topic_metrics
         return {}
     except Exception as e:
         logger.error("[ERR-312] Failed to collect aiokafka topic metrics: {}", e)
@@ -283,7 +302,7 @@ def collect_all_metrics(producer: Any, library: str, tracker_uuid: str = None) -
             topic_metrics = collect_confluent_topic_metrics(tracker_uuid)
             node_metrics = collect_confluent_node_metrics(tracker_uuid)
         elif library == "aiokafka":
-            producer_metrics = collect_aiokafka_metrics(producer)
+            producer_metrics = collect_aiokafka_producer_metrics(producer)
             topic_metrics = collect_aiokafka_topic_metrics(producer)
             node_metrics = collect_aiokafka_node_metrics(producer)
         else: