import json from typing import List from redis import Redis from time import time import ingestor.metrics as metrics from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.utils.logger import Logger from sientia_do.temporal.activities.base import BaseActivity class ResourceManager(BaseActivity): def __init__( self, host: str, port: int, lease_ttl: int, heartbeat_ttl: int, pod_id: str, metadata: dict, logger: Logger, notification_handler: NotificationHandler, username: str | None = None, password: str | None = None, ) -> None: self.pod_id = pod_id try: self.redis = Redis( host=host, port=port, decode_responses=True, username=username, password=password, ) self.redis.ping() metrics.REDIS_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(1) except Exception as e: logger.error(f"Failed to connect to Redis: {e}") metrics.REDIS_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0) raise self.lease_ttl = lease_ttl self.heartbeat_ttl = heartbeat_ttl self.metadata = metadata BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler, set_error_counter_metric=False) def _execute_redis_op(self, operation_name: str, func, *args, **kwargs): """Wrapper to execute Redis operations and record metrics.""" start_time = time() try: result = func(*args, **kwargs) metrics.REDIS_OPERATIONS_TOTAL.labels( pod_id=self.pod_id, operation=operation_name ).inc() duration = time() - start_time metrics.REDIS_OPERATIONS_DURATION.labels( pod_id=self.pod_id, operation=operation_name ).observe(duration) return result except Exception as e: metrics.REDIS_OPERATIONS_ERRORS.labels( pod_id=self.pod_id, operation=operation_name ).inc() self.send_notification( metadata=self.metadata, notification_id=f"REDIS_OPERATION_ERROR_{operation_name}", message=f"Error in Redis operation '{operation_name}': {e}", block="redis_manager", level=NotificationLevel.ERROR, ) raise def get(self, key: str) -> dict: """ Retrieve a value from Redis by its key and return it as a dictionary. Args: key (str): The key to look up in Redis. Returns: dict: The value associated with the key, parsed as a dictionary, or None if the key does not exist or the value is empty. """ history = self._execute_redis_op("get", self.redis.get, key) return json.loads(history) if history else None def get_tag_slot(self, id: str) -> dict: """ Retrieve the tag slot information for a given ID. Args: id (str): The unique identifier of the tag slot to retrieve. Returns: dict: A dictionary containing the tag slot information associated with the given ID. """ return self.get(f"slot:opc_tags:{id}") def ingestor_heartbeat(self) -> None: """ Sends a heartbeat signal to Redis to indicate that the ingestor is active. This method sets a key in Redis with a specific format that includes the ingestor's pod ID. The key is set with a value of 1 and an expiration time defined by `self.heartbeat_ttl`. This allows monitoring systems to track the activity and health of the ingestor. Returns: None """ self._execute_redis_op( "set", self.redis.set, f"heartbeat:ingestor:{self.pod_id}", 1, ex=self.heartbeat_ttl, ) def lease_tag(self, tag_id: str) -> bool: """ Attempts to lease a tag by setting a key in Redis with a specified TTL (time-to-live). This method uses the Redis `SET` command with the `NX` option to ensure that the key is only set if it does not already exist. The key is set with an expiration time defined by `lease_ttl`. Args: tag_id (str): The unique identifier of the tag to be leased. Returns: bool: True if the lease was successfully acquired, False otherwise. """ return self._execute_redis_op( "set_nx", self.redis.set, f"lease:opc_tags:{tag_id}", self.pod_id, nx=True, ex=self.lease_ttl, ) def renew_tag_lease(self, tag_id: str) -> bool: """ Renews the lease for a specific OPC tag if the current pod holds the lease. This method checks if the current pod (identified by `self.pod_id`) holds the lease for the given OPC tag. If so, it extends the lease by resetting its expiration time in Redis to the configured lease TTL (`self.lease_ttl`). Args: tag_id (str): The identifier of the OPC tag whose lease is to be renewed. Returns: bool: True if the lease was successfully renewed, False otherwise. """ current = self._execute_redis_op( "get", self.redis.get, f"lease:opc_tags:{tag_id}" ) if current == self.pod_id: self._execute_redis_op( "expire", self.redis.expire, f"lease:opc_tags:{tag_id}", self.lease_ttl ) return True return False def drop_tag_lease(self, tag_id: str) -> None: """ Drops the lease for a specific OPC tag. This method removes the lease for the given OPC tag by deleting the corresponding key in Redis. Args: tag_id (str): The identifier of the OPC tag whose lease is to be dropped. Returns: None """ self._execute_redis_op("delete", self.redis.delete, f"lease:opc_tags:{tag_id}") def get_all_ingestors(self) -> List[str]: """ Retrieves all active ingestors from Redis. This method fetches all keys in Redis that match the pattern for ingestor leases and returns a list of active ingestors. Returns: list: A list of active ingestors. """ return self._execute_redis_op("keys", self.redis.keys, "heartbeat:ingestor:*") def get_all_slots(self) -> List[str]: """ Retrieves the number of slots available in Redis. This method counts the number of keys in Redis that match the pattern for OPC tag leases and returns the count. Returns: int: The number of slots available. """ return self._execute_redis_op("keys", self.redis.keys, "slot:opc_tags:*") def get_all_leases(self) -> List[str]: """ Retrieves all active leases from Redis. This method fetches all keys in Redis that match the pattern for OPC tag leases and returns a list of active leases. Returns: list: A list of active leases. """ return self._execute_redis_op("keys", self.redis.keys, "lease:opc_tags:*")