Python Prometheus client: 高效管理与获取度量指标对象

from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram, Summary, Enum, write_to_textfile
from typing import Dict, Union

# 定义所有可能的度量指标类型
MetricType = Union[Counter, Gauge, Histogram, Summary, Enum]

class PrometheusMetricsManager:
    def __init__(self):
        self._registry = CollectorRegistry()
        self._metrics: Dict[str, MetricType] = {}

    def get_registry(self) -> CollectorRegistry:
        """获取内部的CollectorRegistry实例。"""
        return self._registry

    def register_metric(self, metric: MetricType):
        """
        注册单个度量指标到注册表并存储在管理器中。
        """
        # 注册到Prometheus的CollectorRegistry
        self._registry.register(metric)
        # 存储到自定义管理器中，以便后续获取
        # 注意：这里我们假设metric.name在Prometheus客户端中是唯一的
        # 对于带有标签的指标，name是基础名称，实际存储的可能是MetricWithLabels
        # 为了简化，我们直接使用metric.name作为key
        # 如果需要区分带标签和不带标签的同名指标，需要更复杂的键策略
        if hasattr(metric, '_name'): # 对于Counter, Gauge等，直接访问_name
             self._metrics[metric._name] = metric
        else: # 对于其他可能没有直接_name属性的复杂指标，需要根据其描述获取名称
             # 这是一个简化的处理，实际应用可能需要更健壮的逻辑
             # 例如，通过metric.describe()获取MetricFamilySamples，再提取name
             print(f"Warning: Metric {metric} might not have a direct '_name' attribute. Using fallback.")
             # 尝试从describe()获取第一个样本的名称
             try:
                 metric_name = next(iter(metric.describe())).name
                 self._metrics[metric_name] = metric
             except Exception:
                 print(f"Could not determine name for metric: {metric}")


    def get_metric(self, name: str) -> MetricType | None:
        """
        根据名称获取已注册的度量指标对象。
        """
        return self._metrics.get(name)

# 使用示例
if __name__ == "__main__":
    manager = PrometheusMetricsManager()

    # 创建并注册Counter
    request_counter = Counter("http_requests_total", "Total HTTP requests.", registry=manager.get_registry())
    manager.register_metric(request_counter)

    # 创建并注册Gauge
    in_progress_gauge = Gauge("http_requests_in_progress", "HTTP requests in progress.", registry=manager.get_registry())
    manager.register_metric(in_progress_gauge)

    # 模拟操作
    request_counter.inc(5)
    in_progress_gauge.set(2)

    # 从管理器中获取Counter并继续操作
    retrieved_counter = manager.get_metric("http_requests_total")
    if retrieved_counter and isinstance(retrieved_counter, Counter):
        retrieved_counter.inc(3)
        print(f"Incremented http_requests_total to: {retrieved_counter._value}") # 直接访问私有属性查看值

    # 从管理器中获取Gauge并继续操作
    retrieved_gauge = manager.get_metric("http_requests_in_progress")
    if retrieved_gauge and isinstance(retrieved_gauge, Gauge):
        retrieved_gauge.set(1)
        print(f"Set http_requests_in_progress to: {retrieved_gauge._value}")

    # 将指标写入文件以验证
    write_to_textfile("metrics_output_manager.prom", manager.get_registry())
    print("Metrics written to metrics_output_manager.prom")

from prometheus_client import CollectorRegistry, Counter, Gauge, write_to_textfile
from prometheus_client.registry import Collector # Collector是所有指标的基类
from typing import Optional

class CustomCollectorRegistry(CollectorRegistry):
    def get_metric(self, name: str) -> Optional[Collector]:
        """
        线程安全地从注册表中获取已注册的度量指标对象。
        """
        with self._lock: # 使用内部锁保证线程安全
            # _names_to_collectors 是一个内部字典，存储了所有已注册的指标
            # 键是指标的名称（不含_total, _bucket等后缀），值是指标对象本身
            return self._names_to_collectors.get(name)

# 使用示例
if __name__ == "__main__":
    # 创建自定义注册表实例
    custom_registry = CustomCollectorRegistry()

    # 创建并注册Counter
    my_counter = Counter("my_app_requests_total", "Total requests for my application.", registry=custom_registry)
    my_counter.inc(10)

    # 创建并注册Gauge
    my_gauge = Gauge("my_app_current_users", "Current active users.", registry=custom_registry)
    my_gauge.set(5)

    # 从自定义注册表中获取Counter并操作
    retrieved_counter_obj = custom_registry.get_metric("my_app_requests_total")
    if retrieved_counter_obj and isinstance(retrieved_counter_obj, Counter):
        retrieved_counter_obj.inc(7)
        print(f"Incremented my_app_requests_total to: {retrieved_counter_obj._value}")

    # 从自定义注册表中获取Gauge并操作
    retrieved_gauge_obj = custom_registry.get_metric("my_app_current_users")
    if retrieved_gauge_obj and isinstance(retrieved_gauge_obj, Gauge):
        retrieved_gauge_obj.set(8)
        print(f"Set my_app_current_users to: {retrieved_gauge_obj._value}")

    # 尝试获取一个不存在的指标
    non_existent_metric = custom_registry.get_metric("non_existent_metric")
    if non_existent_metric is None:
        print("Successfully handled non-existent metric retrieval.")

    # 将指标写入文件以验证
    write_to_textfile("metrics_output_custom_registry.prom", custom_registry)
    print("Metrics written to metrics_output_custom_registry.prom")