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

import asyncio
from os import getenv
from copy import deepcopy
from typing import Dict, Any

from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger

from ingestor.managers.ingestor_manager import IngestorManager

import ingestor.metrics as metrics


class Ingestor:
    def __init__(self):
        """
        Initializes the ingestor with configuration values retrieved from environment variables.
        Environment Variables:
            KAFKA_SERVERS (str): Comma-separated list of Kafka server addresses.
                Defaults to "localhost:9092".
            REDIS_HOST (str): Hostname of the Redis server. Defaults to "localhost".
            REDIS_PORT (int): Port number of the Redis server. Defaults to 6379.
            LEASE_TTL (int): Time-to-live for leases in seconds. Defaults to 10.
            HEARTBEAT_TTL (int): Time-to-live for heartbeats in seconds. Defaults to 20.
            HOSTNAME (str): Identifier for the current pod or host. Defaults to "localhost".
            POLL_INTERVAL (int): Interval in seconds for polling operations. Defaults to 5.
            MONGODB_URL (str): URL of the MongoDB server. Defaults to "localhost:27017".
            MONGODB_USERNAME (str): Username for the MongoDB server. Defaults to "sientia".
            MONGODB_PASSWORD (str): Password for the MongoDB server. Defaults to "sientia".
            MONGODB_DATABASE (str): Name of the MongoDB database. Defaults to "sientia".
        Attributes:
            kafka_servers (list): List of Kafka server addresses.
            redis_host (str): Hostname of the Redis server.
            redis_port (int): Port number of the Redis server.
            lease_ttl (int): Time-to-live for leases in seconds.
            heartbeat_ttl (int): Time-to-live for heartbeats in seconds.
            pod_id (str): Identifier for the current pod or host.
            poll_interval (int): Interval in seconds for polling operations.
        """

        kafka_servers = getenv("KAFKA_SERVERS", "localhost:9092")
        export_to_kafka = getenv("EXPORT_TO_KAFKA", "false")

        if export_to_kafka and export_to_kafka == "true":
            export_to_kafka = True
        else:
            export_to_kafka = False

        self.export_to_kafka = export_to_kafka
        self.redis_host = getenv("REDIS_HOST", "localhost")
        self.redis_port = int(getenv("REDIS_PORT", "6379"))
        self.redis_username = getenv("REDIS_USERNAME", None)
        self.redis_password = getenv("REDIS_PASSWORD", None)
        self.lease_ttl = int(getenv("LEASE_TTL", "10"))
        self.heartbeat_ttl = int(getenv("HEARTBEAT_TTL", "20"))
        self.pod_id = getenv("HOSTNAME", "localhost")
        self.poll_interval = int(getenv("POLL_INTERVAL", "5"))
        mongo_url = getenv("MONGODB_URL", "localhost:27017")
        mongo_username = getenv("MONGODB_USERNAME", "sientia")
        mongo_password = getenv("MONGODB_PASSWORD", "sientia")
        self.mongo_database = getenv("MONGODB_DATABASE", "sientia")
        self.mongo_connection_string = f"mongodb://{mongo_username}:{mongo_password}@{mongo_url}"

        self.kafka_servers = kafka_servers.split(",")
        self.logger = get_logger(__name__)
        self.notification_handler = NotificationHandler(
            connection_string=self.mongo_connection_string,
            database=self.mongo_database,
            logger=self.logger,
            project_name="opc_ingestor"
        )

        self.metadata = {
            'model_id': '-',
            'model_name': '-',
            'workflow_name': 'opc_ingestor',
            'schema_name': 'opc_ingestor',
            'pod_id': self.pod_id,
        }
        self.ingestor_manager = None

    async def shutdown(self):
        if self.ingestor_manager:
            await self.ingestor_manager.shutdown()

    async def handle_acquired_tags(self, acquired):
        """
        Handles the acquired tags by subscribing to them if available.
        This method checks if there are any acquired tags. If no tags are acquired,
        it logs a warning indicating that no slots are available. Otherwise, it
        updates the OPC servers and subscribes to the acquired tags.
        Args:
            acquired (list): A list of acquired tags to be processed. If the list
                             is empty or None, no action is taken other than logging
                             a warning.
        """

        if not acquired:
            self.logger.warning("No slots available")

        else:
            # Subscribe to acquired slots
            self.ingestor_manager.update_opc_servers()
            await self.ingestor_manager.subscribe_to_tags(acquired)

    async def prepare_ingestor(self):
        """
        Prepares the ingestor by initializing the IngestorManager, declaring the ingestor as active,
        acquiring slot leases, and handling the acquired tags.
        This method performs the following steps:
        1. Initializes the `IngestorManager` with the necessary configuration parameters.
        2. Declares the ingestor as active by calling `declare_active` on the `IngestorManager`.
        3. Acquires slot leases using the `get_slot_leases` method of the `IngestorManager`.
        4. Logs the acquired slots and processes them using the `handle_acquired_tags` method.
        Attributes:
            self.kafka_servers (list): List of Kafka server addresses.
            self.redis_host (str): Redis server hostname.
            self.redis_port (int): Redis server port.
            self.lease_ttl (int): Time-to-live for slot leases.
            self.heartbeat_ttl (int): Time-to-live for heartbeat signals.
            self.pod_id (str): Identifier for the current pod.
            self.poll_interval (int): Interval for polling operations.
            self.logger (Logger): Logger instance for logging messages.
        Raises:
            Exception: If any error occurs during the initialization or lease acquisition process.
        """

        self.ingestor_manager = IngestorManager(
            kafka_servers=self.kafka_servers,
            redis_data={
                'host': self.redis_host,
                'port': self.redis_port,
                'username': self.redis_username,
                'password': self.redis_password,
            },
            lease_ttl=self.lease_ttl,
            heartbeat_ttl=self.heartbeat_ttl,
            poll_interval=self.poll_interval,
            mongo_connection_string=self.mongo_connection_string,
            mongo_database=self.mongo_database,
            metadata=self.metadata,
            logger=self.logger,
            notification_handler=self.notification_handler,
            export_to_kafka=self.export_to_kafka,
        )

        # Declare ingestor ative
        self.ingestor_manager.declare_active()

        # Get slot lease
        acquired = self.ingestor_manager.get_slot_leases()
        self.logger.info(f"Acquired slots: {acquired}")

        await self.handle_acquired_tags(acquired)

        metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(
            len(self.ingestor_manager.managed_tags)
        )  # Set initial

    def manage_no_slots(self, number_of_slots: int):
        """
        Manages the scenario where there are no slots assigned to the ingestor.
        This method checks if the ingestor is active (i.e., has no managed tags)
        and if the number of available slots is greater than zero. If both
        conditions are met, it attempts to acquire a slot lease and handles
        the acquired tags accordingly.
        Args:
            number_of_slots (int): The number of available slots.
        """

        if not self.ingestor_manager.managed_tags and number_of_slots > 0:
            # This ingestor is active and has no slots, so we need to try to

            # Get slot lease
            self.ingestor_manager.get_slot_leases(1)

    async def manage_leases(
        self, available_slots: int, lacking_ingestors: int, slot_diff: int
    ):
        """
        Manages the allocation and deallocation of slot leases for ingestors based on
        the number of available slots, lacking ingestors, and slot differences.
        Args:
            available_slots (int): The number of slots currently available for allocation.
            lacking_ingestors (int): The number of ingestors that are active and without slots.
            slot_diff (int): The difference between the total slots and the required slots.
        Behavior:
            - If there are available slots and lacking ingestors, attempts to acquire slot leases
              for the available slots and processes the acquired tags.
            - If there are no lacking ingestors but there are extra slots (slot_diff > 0),
              releases the extra slot leases to ensure proper allocation.
        Logs:
            - Logs the number of available slots when attempting to acquire leases.
            - Logs the number of extra slots when releasing leases.
        """

        if available_slots > 0 and lacking_ingestors > 0:
            # Some ingestors are innactive, so theres "available_slots" slots available
            self.logger.info(f"Slots available: {available_slots}")

            # Get slot lease
            self.ingestor_manager.get_slot_leases(available_slots)

        elif lacking_ingestors <= 0 and slot_diff > 0:

            self.logger.info(f"Extra slots available: {slot_diff}")
            # There's enough slots for all ingestors, but this ingestor has more than one slot
            # So we need to drop the extra leases

            overleases = list(self.ingestor_manager.managed_tags.keys())[1:]

            self.ingestor_manager.drop_slot_leases(overleases)

            for lease in overleases:
                await self.ingestor_manager.unsubscribe_slot(lease)
                self.ingestor_manager.managed_tags.pop(lease)

            # Update metric after removal
            metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(
                len(self.ingestor_manager.managed_tags)
            )

    async def update_ingestor_manager(self, old_managed_tags: Dict[str, Any]):
        """
        Updates the ingestor manager with the new managed tags.
        Args:
            old_managed_tags (Dict[str, Any]): The old managed tags.
        """

        self.logger.debug(
            f"Current managed tags: {self.ingestor_manager.managed_tags}")

        await self.ingestor_manager.update_opc_servers()
        new_managed_tags = deepcopy(self.ingestor_manager.managed_tags)

        self.logger.debug(
            f"Comparing new managed tags {new_managed_tags} with old managed tags {old_managed_tags}"
        )

        keys = set(new_managed_tags) | set(old_managed_tags)

        changes = {k: (new_managed_tags.get(k), old_managed_tags.get(k))
                   for k in keys if new_managed_tags.get(k) != old_managed_tags.get(k)}

        self.logger.debug(f"Changes: {changes}")

        for slot, config in new_managed_tags.items():
            if slot not in old_managed_tags:
                self.logger.info(f"Subscribing to new slot {slot}")
                await self.ingestor_manager.subscribe_to_tags({slot: config})
                continue

            if config != old_managed_tags[slot]:
                self.logger.info(f"Resubscribing to slot {slot}")
                await self.ingestor_manager.unsubscribe_slot(slot)
                await self.ingestor_manager.subscribe_to_tags({slot: config})

        for slot in old_managed_tags.keys():
            if slot not in new_managed_tags:
                self.logger.info(f"Unsubscribing from slot {slot}")
                await self.ingestor_manager.unsubscribe_slot(slot)

        # Ensure the gauge is updated after any potential changes here
        metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(
            len(self.ingestor_manager.managed_tags)
        )

    async def loop(self):
        """
        Executes the main loop for managing ingestors and slots.
        This method performs the following tasks:
        1. Declares the ingestor as active.
        2. Logs the start of the polling process for slot updates.
        3. Retrieves the list of active ingestors and the number of available slots.
        4. Handles scenarios where no slots are available.
        5. Calculates the difference between the number of slots and active ingestors,
           as well as the difference in managed tags.
        6. Manages leases based on the calculated differences.
        7. Logs the current state of active ingestors, slots, managed tags, and servers.
        8. Logs a message if no slots are acquired during the loop.
        9. Updates the configuration of OPC servers.
        This method is intended to be called repeatedly to ensure the ingestor
        manager operates correctly and maintains synchronization with the slots
        and OPC servers.
        """

        self.ingestor_manager.declare_active()

        self.logger.info("Polling for slot updates...")
        # Get active ingestors

        current_managed_tags = deepcopy(self.ingestor_manager.managed_tags)

        ingestors = self.ingestor_manager.get_active_ingestors()
        number_of_ingestors = len(ingestors)
        number_of_leases = self.ingestor_manager.get_number_of_leases()
        number_of_slots = self.ingestor_manager.get_number_of_slots()

        # Update active ingestors gauge
        metrics.ACTIVE_INGESTORS.set(number_of_ingestors)

        # Handle no slots
        self.logger.info("Managing no slots...")
        self.manage_no_slots(number_of_slots)

        available_slots = number_of_slots - number_of_leases
        lacking_ingestors = number_of_slots - number_of_ingestors
        slot_diff = len(self.ingestor_manager.managed_tags) - 1

        self.logger.info("Managing leases...")
        await self.manage_leases(available_slots, lacking_ingestors, slot_diff)

        # Update managed slots gauge
        metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(
            len(self.ingestor_manager.managed_tags)
        )

        self.logger.debug(
            f"Active ingestors: {ingestors}, "
            f"Number of slots: {number_of_slots}, "
            f"Number of leases: {number_of_leases}, "
            f"Managed tags: {self.ingestor_manager.managed_tags}, "
            f"Managed servers: {self.ingestor_manager.opc_managers}"
        )
        if not self.ingestor_manager.managed_tags:
            # No slots acquired
            self.logger.info("No slots acquired in this loop")

        # Update opc servers
        self.logger.info("Updating slot config...")
        self.ingestor_manager.update_slot_config()

        # Check OPC cycles
        self.logger.info("Checking OPC servers integrity...")
        self.ingestor_manager.check_opc_servers_integrity()

        self.logger.info("Updating managed tags...")
        await self.update_ingestor_manager(current_managed_tags)