support of metrics collection for Confluent lib

Author: idanasulin2706

examples/requirements.txt

@@ -1,4 +1,4 @@
 kafka-python==2.2.14
-confluent-kafka==2.3.0
+confluent-kafka==2.11.0
 aiokafka==0.10.0 
 aws-msk-iam-sasl-signer-python==1.0.2

superclient/agent/interceptor.py

@@ -1,10 +1,12 @@
 """Producer interception functionality."""
 
 import os
+import uuid
 from typing import Any, Dict
 
 from ..util.logger import get_logger
 from ..util.config import get_topics_list, is_disabled
+from ..util.metrics import configure_confluent_stats_callback
 from .metadata import fetch_metadata_sync, optimal_cfg, _DEFAULTS
 from ..core.reporter import send_clients_msg
 from ..core.manager import normalize_bootstrap
@@ -345,10 +347,18 @@ def __init__(self, conf: Dict[str, Any], *args, **kwargs):
                             logger.debug("Overriding configuration: {} ((not set) -> {})", k, v)
                         conf[k] = v
 
+                
+                # Generate UUID for this producer
+                tracker_uuid = str(uuid.uuid4())
+                
+                # Configure stats callback for metrics collection
+                conf = configure_confluent_stats_callback(conf, tracker_uuid)
+                
                 # Create the producer with optimized configuration
                 self._producer = Producer(conf, *args, **kwargs)
 
                 report_interval = metadata.get("report_interval_ms") if metadata else _DEFAULT_REPORT_INTERVAL_MS
+                # Create tracker with the generated UUID
                 self._tracker = ProducerTracker(
                     lib="confluent",
                     producer=self._producer,
@@ -357,10 +367,12 @@ def __init__(self, conf: Dict[str, Any], *args, **kwargs):
                     orig_cfg=orig_cfg,
                     opt_cfg=opt_cfg,
                     report_interval_ms=int(report_interval or _DEFAULT_REPORT_INTERVAL_MS),
-                    error=error_msg,  # Store error message in tracker
+                    error=error_msg,
                     metadata=metadata,
                     topics_env=topics_env,
+                    uuid=tracker_uuid,  # Use the generated UUID
                 )
+                
                 Heartbeat.register_tracker(self._tracker)
 
                 send_clients_msg(self._tracker, error_msg)
@@ -397,12 +409,15 @@ def __del__(self):
                     self._superstream_closed = True
                     self._tracker.close()
                     Heartbeat.unregister_tracker(self._tracker.uuid)
+                    
+                    # Remove metrics extractor from registry
+                    from ..util.metrics import remove_producer_metrics_extractor
+                    remove_producer_metrics_extractor(self._tracker.uuid)
+                    
                     logger.debug("Superstream tracking stopped for confluent-kafka producer with client_id: {}", 
                                getattr(self._tracker, 'client_id', 'unknown'))
                 except Exception as e:
                     logger.error("Error during automatic cleanup: {}", e)
-            else:
-                logger.debug("Producer already cleaned up or no tracker found")
 
         def __getattr__(self, name):
             """Delegate all other attributes to the underlying producer."""

superclient/agent/tracker.py

@@ -24,8 +24,9 @@ def __init__(
         error: str = "",
         metadata: Optional[Dict[str, Any]] = None,
         topics_env: Optional[list[str]] = None,
+        uuid: Optional[str] = None,
     ) -> None:
-        self.uuid = str(uuid.uuid4())
+        self.uuid = uuid if uuid else str(uuid.uuid4())
         self.library = lib
         self.producer = producer
         self.bootstrap = bootstrap
@@ -95,7 +96,7 @@ def _update_cache_from_producer(self):
             from ..util.metrics import collect_all_metrics, merge_metrics_with_cache
 
             # Collect fresh metrics
-            current_producer_metrics, current_topic_metrics, current_node_metrics = collect_all_metrics(self.producer, self.library)
+            current_producer_metrics, current_topic_metrics, current_node_metrics = collect_all_metrics(self.producer, self.library, self.uuid)
 
             # Get current cached metrics
             cached_producer_metrics, cached_topic_metrics, cached_node_metrics = self.get_cached_metrics()