respond to comments

Signed-off-by: Tim Li <ltim@uber.com>

Author: timl3136

cadence/_internal/visibility/__init__.py

@@ -1,12 +1,12 @@
 """Visibility and metrics collection components for Cadence client."""
 
-from .metrics import MetricsHandler, NoOpMetricsHandler, get_default_handler
-from .prometheus import PrometheusMetrics, PrometheusConfig
+from .metrics import MetricsEmitter, NoOpMetricsEmitter
+from .prometheus import PrometheusMetrics, PrometheusConfig, CadenceMetrics
 
 __all__ = [
-    "MetricsHandler",
-    "NoOpMetricsHandler", 
-    "get_default_handler",
+    "MetricsEmitter",
+    "NoOpMetricsEmitter", 
     "PrometheusMetrics", 
     "PrometheusConfig",
+    "CadenceMetrics",
 ]

cadence/_internal/visibility/metrics.py

@@ -13,10 +13,9 @@ class MetricType(Enum):
     COUNTER = "counter"
     GAUGE = "gauge"
     HISTOGRAM = "histogram"
-    TIMER = "timer"
 
 
-class MetricsHandler(Protocol):
+class MetricsEmitter(Protocol):
     """Protocol for metrics collection backends."""
 
     def counter(
@@ -31,11 +30,6 @@ def gauge(
         """Send a gauge metric."""
         ...
 
-    def timer(
-        self, key: str, duration: float, tags: Optional[Dict[str, str]] = None
-    ) -> None:
-        """Send a timer metric."""
-        ...
 
     def histogram(
         self, key: str, value: float, tags: Optional[Dict[str, str]] = None
@@ -44,8 +38,8 @@ def histogram(
         ...
 
 
-class NoOpMetricsHandler:
-    """No-op metrics handler that discards all metrics."""
+class NoOpMetricsEmitter:
+    """No-op metrics emitter that discards all metrics."""
 
     def counter(
         self, key: str, n: int = 1, tags: Optional[Dict[str, str]] = None
@@ -57,30 +51,10 @@ def gauge(
     ) -> None:
         pass
 
-    def timer(
-        self, key: str, duration: float, tags: Optional[Dict[str, str]] = None
-    ) -> None:
-        pass
 
     def histogram(
         self, key: str, value: float, tags: Optional[Dict[str, str]] = None
     ) -> None:
         pass
 
 
-# Global default handler
-_default_handler: Optional[MetricsHandler] = None
-
-
-def get_default_handler() -> MetricsHandler:
-    """Get the default global metrics handler."""
-    global _default_handler
-    if _default_handler is None:
-        _default_handler = NoOpMetricsHandler()
-    return _default_handler
-
-
-def set_default_handler(handler: MetricsHandler) -> None:
-    """Set the default global metrics handler."""
-    global _default_handler
-    _default_handler = handler

cadence/_internal/visibility/prometheus.py

@@ -2,7 +2,8 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Dict, Optional, Any
+from enum import Enum
+from typing import Dict, Optional
 
 from prometheus_client import (  # type: ignore[import-not-found]
     REGISTRY,
@@ -11,10 +12,10 @@
     Gauge,
     Histogram,
     generate_latest,
-    push_to_gateway,
-    start_http_server,
 )
 
+from cadence._internal.visibility.metrics import MetricsEmitter
+
 
 logger = logging.getLogger(__name__)
 
@@ -23,16 +24,6 @@
 class PrometheusConfig:
     """Configuration for Prometheus metrics."""
 
-    # HTTP server configuration
-    enable_http_server: bool = False
-    http_port: int = 8000
-    http_addr: str = "0.0.0.0"
-
-    # Push gateway configuration
-    enable_push_gateway: bool = False
-    push_gateway_url: str = "localhost:9091"
-    push_job_name: str = "cadence_client"
-
     # Metric name prefix
     metric_prefix: str = "cadence_"
 
@@ -43,7 +34,7 @@ class PrometheusConfig:
     registry: Optional[CollectorRegistry] = None
 
 
-class PrometheusMetrics:
+class PrometheusMetrics(MetricsEmitter):
     """Prometheus metrics collector implementation."""
 
     def __init__(self, config: Optional[PrometheusConfig] = None):
@@ -55,12 +46,6 @@ def __init__(self, config: Optional[PrometheusConfig] = None):
         self._gauges: Dict[str, Gauge] = {}
         self._histograms: Dict[str, Histogram] = {}
 
-        # HTTP server handle
-        self._http_server: Optional[Any] = None
-
-        if self.config.enable_http_server:
-            self.start_http_server()
-
     def _get_metric_name(self, name: str) -> str:
         """Get the full metric name with prefix."""
         return f"{self.config.metric_prefix}{name}"
@@ -126,23 +111,6 @@ def _get_or_create_histogram(
 
         return self._histograms[metric_name]
 
-    def _get_or_create_timer(
-        self, name: str, labels: Optional[Dict[str, str]]
-    ) -> Histogram:
-        """Get or create a timer metric (implemented as histogram)."""
-        metric_name = self._get_metric_name(name)
-
-        if metric_name not in self._histograms:
-            label_names = list(self._merge_labels(labels).keys()) if labels else []
-            self._histograms[metric_name] = Histogram(
-                metric_name,
-                f"Timer metric for {name}",
-                labelnames=label_names,
-                registry=self.registry,
-            )
-            logger.debug(f"Created timer metric: {metric_name}")
-
-        return self._histograms[metric_name]
 
     def counter(
         self, key: str, n: int = 1, tags: Optional[Dict[str, str]] = None
@@ -176,21 +144,6 @@ def gauge(
         except Exception as e:
             logger.error(f"Failed to send gauge {key}: {e}")
 
-    def timer(
-        self, key: str, duration: float, tags: Optional[Dict[str, str]] = None
-    ) -> None:
-        """Send a timer metric - implemented as histogram."""
-        try:
-            timer = self._get_or_create_timer(key, tags)
-            merged_tags = self._merge_labels(tags)
-
-            if merged_tags:
-                timer.labels(**merged_tags).observe(duration)
-            else:
-                timer.observe(duration)
-
-        except Exception as e:
-            logger.error(f"Failed to send timer {key}: {e}")
 
     def histogram(
         self, key: str, value: float, tags: Optional[Dict[str, str]] = None
@@ -208,44 +161,6 @@ def histogram(
         except Exception as e:
             logger.error(f"Failed to send histogram {key}: {e}")
 
-    def start_http_server(self) -> None:
-        """Start HTTP server to expose metrics."""
-        if self._http_server is not None:
-            logger.warning("HTTP server already started")
-            return
-
-        try:
-            server_result = start_http_server(
-                self.config.http_port,
-                addr=self.config.http_addr,
-                registry=self.registry,
-            )
-            self._http_server = server_result
-            logger.info(
-                f"Prometheus metrics HTTP server started on "
-                f"{self.config.http_addr}:{self.config.http_port}"
-            )
-        except Exception as e:
-            logger.error(f"Failed to start HTTP server: {e}")
-            raise
-
-    def push_to_gateway(self) -> None:
-        """Push metrics to Prometheus Push Gateway."""
-        if not self.config.enable_push_gateway:
-            logger.warning("Push gateway not enabled")
-            return
-
-        try:
-            push_to_gateway(
-                self.config.push_gateway_url,
-                job=self.config.push_job_name,
-                registry=self.registry,
-            )
-            logger.debug(f"Pushed metrics to gateway: {self.config.push_gateway_url}")
-        except Exception as e:
-            logger.error(f"Failed to push to gateway: {e}")
-            raise
-
     def get_metrics_text(self) -> str:
         """Get metrics in Prometheus text format."""
         try:
@@ -255,19 +170,9 @@ def get_metrics_text(self) -> str:
             logger.error(f"Failed to generate metrics text: {e}")
             return ""
 
-    def shutdown(self) -> None:
-        """Shutdown the metrics collector."""
-        if self._http_server:
-            try:
-                self._http_server.shutdown()
-                self._http_server = None
-                logger.info("Prometheus HTTP server shutdown")
-            except Exception as e:
-                logger.error(f"Failed to shutdown HTTP server: {e}")
-
 
 # Default Cadence metrics names
-class CadenceMetrics:
+class CadenceMetrics(Enum):
     """Standard Cadence client metrics."""
 
     # Workflow metrics

cadence/client.py

@@ -12,6 +12,7 @@
 from grpc.aio import Channel, ClientInterceptor, secure_channel, insecure_channel
 from cadence.api.v1.service_workflow_pb2_grpc import WorkflowAPIStub
 from cadence.data_converter import DataConverter, DefaultDataConverter
+from cadence._internal.visibility.metrics import MetricsEmitter, NoOpMetricsEmitter
 
 
 class ClientOptions(TypedDict, total=False):
@@ -24,6 +25,7 @@ class ClientOptions(TypedDict, total=False):
     channel_arguments: dict[str, Any]
     credentials: ChannelCredentials | None
     compression: Compression
+    metrics_emitter: MetricsEmitter
     interceptors: list[ClientInterceptor]
 
 _DEFAULT_OPTIONS: ClientOptions = {
@@ -34,6 +36,7 @@ class ClientOptions(TypedDict, total=False):
     "channel_arguments": {},
     "credentials": None,
     "compression": Compression.NoCompression,
+    "metrics_emitter": NoOpMetricsEmitter(),
     "interceptors": [],
 }
 
@@ -69,6 +72,10 @@ def worker_stub(self) -> WorkerAPIStub:
     def workflow_stub(self) -> WorkflowAPIStub:
         return self._workflow_stub
 
+    @property
+    def metrics_emitter(self) -> MetricsEmitter:
+        return self._options["metrics_emitter"]
+
     async def ready(self) -> None:
         await self._channel.channel_ready()
 

tests/cadence/_internal/visibility/test_metrics.py

@@ -4,82 +4,48 @@
 
 
 from cadence._internal.visibility.metrics import (
-    MetricsHandler,
+    MetricsEmitter,
     MetricType,
-    NoOpMetricsHandler,
-    get_default_handler,
-    set_default_handler,
+    NoOpMetricsEmitter,
 )
 
 
-class TestMetricsHandler:
-    """Test cases for MetricsHandler protocol."""
+class TestMetricsEmitter:
+    """Test cases for MetricsEmitter protocol."""
 
-    def test_noop_handler(self):
-        """Test no-op handler doesn't raise exceptions."""
-        handler = NoOpMetricsHandler()
+    def test_noop_emitter(self):
+        """Test no-op emitter doesn't raise exceptions."""
+        emitter = NoOpMetricsEmitter()
 
         # Should not raise any exceptions
-        handler.counter("test_counter", 1)
-        handler.gauge("test_gauge", 42.0)
-        handler.histogram("test_histogram", 0.5)
-        handler.timer("test_timing", 1.5)
+        emitter.counter("test_counter", 1)
+        emitter.gauge("test_gauge", 42.0)
+        emitter.histogram("test_histogram", 0.5)
 
-    def test_mock_handler(self):
-        """Test mock handler implementation."""
-        mock_handler = Mock(spec=MetricsHandler)
+    def test_mock_emitter(self):
+        """Test mock emitter implementation."""
+        mock_emitter = Mock(spec=MetricsEmitter)
 
         # Test counter
-        mock_handler.counter("test_counter", 2, {"label": "value"})
-        mock_handler.counter.assert_called_once_with(
+        mock_emitter.counter("test_counter", 2, {"label": "value"})
+        mock_emitter.counter.assert_called_once_with(
             "test_counter", 2, {"label": "value"}
         )
 
         # Test gauge
-        mock_handler.gauge("test_gauge", 100.0, {"env": "test"})
-        mock_handler.gauge.assert_called_once_with(
+        mock_emitter.gauge("test_gauge", 100.0, {"env": "test"})
+        mock_emitter.gauge.assert_called_once_with(
             "test_gauge", 100.0, {"env": "test"}
         )
 
-        # Test timer
-        mock_handler.timer("test_timing", 0.75, {"tag": "value"})
-        mock_handler.timer.assert_called_once_with(
-            "test_timing", 0.75, {"tag": "value"}
-        )
 
         # Test histogram
-        mock_handler.histogram("test_histogram", 2.5, {"env": "prod"})
-        mock_handler.histogram.assert_called_once_with(
+        mock_emitter.histogram("test_histogram", 2.5, {"env": "prod"})
+        mock_emitter.histogram.assert_called_once_with(
             "test_histogram", 2.5, {"env": "prod"}
         )
 
 
-class TestDefaultHandler:
-    """Test cases for default handler management."""
-
-    def test_get_default_handler(self):
-        """Test getting the default handler."""
-        handler = get_default_handler()
-        assert isinstance(handler, NoOpMetricsHandler)
-
-        # Should return the same instance
-        handler2 = get_default_handler()
-        assert handler is handler2
-
-    def test_set_default_handler(self):
-        """Test setting a custom default handler."""
-        original_handler = get_default_handler()
-        custom_handler = NoOpMetricsHandler()
-
-        set_default_handler(custom_handler)
-
-        # Should return the custom handler
-        current_handler = get_default_handler()
-        assert current_handler is custom_handler
-        assert current_handler is not original_handler
-
-        # Restore original for other tests
-        set_default_handler(original_handler)
 
 
 class TestMetricType:
@@ -90,4 +56,3 @@ def test_metric_type_values(self):
         assert MetricType.COUNTER.value == "counter"
         assert MetricType.GAUGE.value == "gauge"
         assert MetricType.HISTOGRAM.value == "histogram"
-        assert MetricType.TIMER.value == "timer"