adds metrics for subscribers

Signed-off-by: Elena Kolevska <[email protected]>

Author: elena-kolevska

metrics.py

@@ -70,10 +70,13 @@ def __init__(self, enable_prometheus: bool = True, prometheus_port: int = 8000,
         self._reconnection_count = 0
         self._reconnection_duration = 0.0
 
-        # OpenTelemetry setup (single collection method)
+        # Setup metrics collection
         if self.enable_otel:
             self._setup_opentelemetry()
 
+        if self.enable_prometheus:
+            self._setup_prometheus_metrics()
+
     def _setup_opentelemetry(self):
         """Setup OpenTelemetry metrics and tracing."""
         try:
@@ -139,7 +142,12 @@ def _setup_opentelemetry(self):
                 unit="1"
             )
 
-
+            # Pub/Sub specific metrics (unified)
+            self.otel_pubsub_operations_counter = self.meter.create_counter(
+                name="redis_pubsub_operations_total",
+                description="Total number of Redis pub/sub operations (publish and receive)",
+                unit="1"
+            )
 
             # Note: Using manual instrumentation instead of automatic Redis instrumentation
             # to have full control over operation labeling and avoid "BATCH" aggregation
@@ -205,6 +213,22 @@ def _setup_prometheus_metrics(self):
             'Average operation latency',
             ['operation'] + base_labels
         )
+
+        # 7. Pub/Sub specific metrics (unified)
+        self.prom_pubsub_operations_total = Counter(
+            'redis_pubsub_operations_total',
+            'Total number of Redis pub/sub operations (publish and receive)',
+            ['channel', 'operation_type', 'subscriber_id', 'status'] + base_labels
+        )
+
+        # Start Prometheus HTTP server
+        try:
+            from prometheus_client import start_http_server
+            start_http_server(self.prometheus_port)
+            self.logger.info(f"Prometheus metrics server started on port {self.prometheus_port}")
+        except Exception as e:
+            self.logger.error(f"Failed to start Prometheus server: {e}")
+            self.enable_prometheus = False
 
 
 
@@ -320,10 +344,40 @@ def update_active_connections(self, count: int):
             }
             self.prom_active_connections.labels(**labels).set(count)
 
+    def record_pubsub_operation(self, channel: str, operation_type: str, subscriber_id: str = None, success: bool = True, error_type: str = None):
+        """Record metrics for a pub/sub operation (publish or receive)."""
 
+        # Update OpenTelemetry metrics
+        if self.enable_otel and hasattr(self, 'otel_pubsub_operations_counter'):
+            labels = {
+                "app_name": self.app_name,
+                "instance_id": self.instance_id,
+                "version": self.version,
+                "run_id": self.run_id,
+                "channel": channel,
+                "operation_type": operation_type,
+                "subscriber_id": subscriber_id or "",
+                "status": "success" if success else "error"
+            }
+            self.otel_pubsub_operations_counter.add(1, labels)
+
+        # Update Prometheus metrics
+        if self.enable_prometheus and hasattr(self, 'prom_pubsub_operations_total'):
+            base_labels = {
+                'app_name': self.app_name,
+                'service_name': self.service_name,
+                'instance_id': self.instance_id
+            }
 
+            status = "success" if success else "error"
+            self.prom_pubsub_operations_total.labels(
+                channel=channel,
+                operation_type=operation_type,
+                subscriber_id=subscriber_id or "",
+                status=status,
+                **base_labels
+            ).inc()
 
-    
     def get_operation_stats(self, operation: str) -> Dict:
         """Get statistics for a specific operation."""
         with self._lock:

redis_client.py

@@ -296,7 +296,19 @@ def llen(self, key: str) -> int:
     # Pub/Sub operations
     def publish(self, channel: str, message: str) -> int:
         """Publish a message to a channel."""
-        return self._execute_with_metrics('PUBLISH', self._client.publish, channel, message)
+        start_time = time.time()
+        try:
+            result = self._client.publish(channel, message)
+            duration = max(0.0, time.time() - start_time)
+            # Record both general operation metrics and pub/sub specific metrics
+            self.metrics.record_operation('PUBLISH', duration, True)
+            self.metrics.record_pubsub_operation(channel, 'PUBLISH', success=True)
+            return result
+        except Exception as e:
+            duration = max(0.0, time.time() - start_time)
+            self.metrics.record_operation('PUBLISH', duration, False, type(e).__name__)
+            self.metrics.record_pubsub_operation(channel, 'PUBLISH', success=False, error_type=type(e).__name__)
+            raise
 
     def pubsub(self):
         """Get a pubsub instance."""

workloads.py

@@ -303,6 +303,9 @@ def __init__(self, config: WorkloadConfig, client: RedisClientManager):
         self._pubsub = None
         self._subscriber_thread = None
         self._stop_subscriber = threading.Event()
+        # Generate unique subscriber ID for this workload instance
+        import uuid
+        self._subscriber_id = f"subscriber_{uuid.uuid4().hex[:8]}"
 
     def _start_subscriber(self):
         """Start subscriber in a separate thread."""
@@ -318,15 +321,27 @@ def _start_subscriber(self):
             while not self._stop_subscriber.is_set():
                 try:
                     # Use get_message with short timeout to be responsive to shutdown
+                    start_time = time.time()
                     message = self._pubsub.get_message(timeout=0.5)
+
                     if message and message['type'] == 'message':
-                        self.logger.debug(f"Received message on {message['channel']}: {message['data']}")
+                        channel_name = message['channel'].decode() if isinstance(message['channel'], bytes) else str(message['channel'])
+
+                        # Record receive metrics using unified pub/sub metric
+                        self.metrics.record_pubsub_operation(channel_name, 'receive', self._subscriber_id, success=True)
+
+                        self.logger.debug(f"Received message on {channel_name}: {message['data']}")
+
                 except (ConnectionError, ValueError) as e:
                     # These are expected during shutdown, break quietly
                     break
                 except Exception as e:
-                    # Other unexpected errors - log only if not shutting down
+                    # Record error metrics if we have channel info
                     if not self._stop_subscriber.is_set():
+                        # Use default channel for error tracking
+                        channels = self.config.get_option("channels", ["test_channel"])
+                        default_channel = channels[0] if channels else "unknown"
+                        self.metrics.record_pubsub_operation(default_channel, 'receive', self._subscriber_id, success=False, error_type=type(e).__name__)
                         self.logger.debug(f"Subscriber error (continuing): {e}")
                     break