sientia-dataops-opc-ingestor/ingestor/managers/resource_manager.py

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)

    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:*")