diff --git a/nucypher/utilities/prometheus/collector.py b/nucypher/utilities/prometheus/collector.py index 88c3ef869..48d629155 100644 --- a/nucypher/utilities/prometheus/collector.py +++ b/nucypher/utilities/prometheus/collector.py @@ -33,10 +33,9 @@ from nucypher.blockchain.eth.registry import BaseContractRegistry from nucypher.datastore.datastore import RecordNotFound from nucypher.datastore.models import Workorder, PolicyArrangement -from prometheus_client.metrics import MetricWrapperBase from prometheus_client.registry import CollectorRegistry -from typing import Dict, List, Union +from typing import Dict, Union ContractAgents = Union[StakingEscrowAgent, WorkLockAgent, PolicyManagerAgent] @@ -432,3 +431,21 @@ class WorkerBondedEventMetricsCollector(EventMetricsCollector): super()._event_occurred(event) self.metrics["current_worker_is_me_gauge"].set( self.contract_agent.get_worker_from_staker(self.staker_address) == self.worker_address) + + +class WorkLockRefundEventMetricsCollector(EventMetricsCollector): + """Collector for WorkLock Refund event.""" + + def __init__(self, staker_address: ChecksumAddress, event_name: str = 'Refund', *args, **kwargs): + super().__init__(event_name=event_name, argument_filters={'sender': staker_address}, *args, **kwargs) + self.staker_address = staker_address + + def initialize(self, metrics_prefix: str, registry: CollectorRegistry) -> None: + super().initialize(metrics_prefix=metrics_prefix, registry=registry) + self.metrics["worklock_deposited_eth_gauge"] = Gauge(f'{metrics_prefix}_worklock_current_deposited_eth', + 'Worklock deposited ETH', + registry=registry) + + def _event_occurred(self, event) -> None: + super()._event_occurred(event) + self.metrics["worklock_deposited_eth_gauge"].set(self.contract_agent.get_deposited_eth(self.staker_address)) diff --git a/nucypher/utilities/prometheus/metrics.py b/nucypher/utilities/prometheus/metrics.py index 7614f5cf6..551fc380a 100644 --- a/nucypher/utilities/prometheus/metrics.py +++ b/nucypher/utilities/prometheus/metrics.py @@ -41,14 +41,15 @@ from nucypher.utilities.prometheus.collector import ( ReStakeEventMetricsCollector, WindDownEventMetricsCollector, WorkerBondedEventMetricsCollector, - CommitmentMadeEventMetricsCollector) + CommitmentMadeEventMetricsCollector, + WorkLockRefundEventMetricsCollector) from typing import List from twisted.internet import reactor, task from twisted.web.resource import Resource -from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent, PolicyManagerAgent +from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent, PolicyManagerAgent, WorkLockAgent class PrometheusMetricsConfig: @@ -205,6 +206,11 @@ def create_metrics_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[Met metrics_prefix=metrics_prefix) collectors.extend(staking_events_collectors) + # WorkLock Events + worklock_events_collectors = create_worklock_events_metric_collectors(ursula=ursula, + metrics_prefix=metrics_prefix) + collectors.extend(worklock_events_collectors) + # Policy Events policy_events_collectors = create_policy_events_metric_collectors(ursula=ursula, metrics_prefix=metrics_prefix) @@ -283,6 +289,25 @@ def create_staking_events_metric_collectors(ursula: 'Ursula', metrics_prefix: st return collectors +def create_worklock_events_metric_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[MetricsCollector]: + """Create collectors for worklock-related events.""" + collectors: List[MetricsCollector] = [] + worklock_agent = ContractAgency.get_agent(WorkLockAgent, registry=ursula.registry) + staker_address = ursula.checksum_address + + # Refund + collectors.append(WorkLockRefundEventMetricsCollector( + event_args_config={ + "refundETH": (Gauge, f'{metrics_prefix}_worklock_refund_refundETH', + 'Refunded ETH'), + }, + staker_address=staker_address, + contract_agent=worklock_agent, + )) + + return collectors + + def create_policy_events_metric_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[MetricsCollector]: """Create collectors for policy-related events.""" collectors: List[MetricsCollector] = []