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:
    """
    Main OPC Ingestor class that orchestrates data collection from OPC UA servers.

    The Ingestor is responsible for:
    - Managing slot leases for load balancing across multiple instances
    - Connecting to and monitoring OPC UA servers
    - Subscribing to OPC tags and collecting real-time data
    - Distributing data to Kafka and MongoDB
    - Providing health monitoring and metrics collection

    The ingestor uses a slot-based architecture where each slot represents
    a collection of OPC tags that can be managed by a single ingestor instance.
    This allows for horizontal scaling and load distribution.

    Environment Variables:
        KAFKA_SERVERS: Comma-separated list of Kafka server addresses (default: "localhost:9092")
        EXPORT_TO_KAFKA: Enable/disable Kafka export (default: "false")
        REDIS_HOST: Redis server hostname (default: "localhost")
        REDIS_PORT: Redis server port (default: 6379)
        REDIS_USERNAME: Redis username (optional)
        REDIS_PASSWORD: Redis password (optional)
        LEASE_TTL: Time-to-live for slot leases in seconds (default: 10)
        HEARTBEAT_TTL: Time-to-live for heartbeats in seconds (default: 20)
        HOSTNAME: Pod identifier (default: "localhost")
        POLL_INTERVAL: Main loop polling interval in seconds (default: 5)
        MONGODB_URL: MongoDB server address (default: "localhost:27017")
        MONGODB_USERNAME: MongoDB username (default: "sientia")
        MONGODB_PASSWORD: MongoDB password (default: "sientia")
        MONGODB_DATABASE: MongoDB database name (default: "sientia")

    Attributes:
        export_to_kafka (bool): Whether to export data to Kafka
        redis_host (str): Redis server hostname
        redis_port (int): Redis server port
        redis_username (str): Redis username (optional)
        redis_password (str): Redis password (optional)
        lease_ttl (int): Time-to-live for slot leases in seconds
        heartbeat_ttl (int): Time-to-live for heartbeat signals in seconds
        pod_id (str): Identifier for the current pod or host
        poll_interval (int): Interval in seconds for polling operations
        mongo_database (str): MongoDB database name
        mongo_connection_string (str): Complete MongoDB connection string
        kafka_servers (list): List of Kafka server addresses
        logger: Logger instance for application logging
        notification_handler: Handler for sending notifications
        metadata (dict): Application metadata for notifications and tracking
        ingestor_manager: Manager instance for coordinating operations
    """

    def __init__(self):
        """
        Initializes the ingestor with configuration values retrieved from environment variables.

        Sets up all necessary connections and configurations for:
        - Kafka connectivity (if enabled)
        - Redis for slot management and coordination
        - MongoDB for data persistence and notifications
        - OPC UA server management
        - Metrics collection and monitoring
        """

        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):
        """
        Gracefully shuts down the ingestor and all its components.

        This method ensures proper cleanup of:
        - OPC UA connections and subscriptions
        - Resource managers and data connections
        - Active slot leases and heartbeats

        Should be called before application termination to prevent resource leaks.
        """
        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 processes the tags that have been allocated to this ingestor
        instance through the slot leasing system. It updates OPC server configurations
        and establishes subscriptions to the allocated 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.

        Behavior:
            - If no tags are acquired, logs a warning about no slots being available
            - If tags are acquired, updates OPC server configurations and subscribes
              to the allocated tags for data collection
        """

        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 all components and acquiring initial slot leases.

        This method performs the following steps:
        1. Initializes the IngestorManager with all necessary configuration parameters
        2. Declares the ingestor as active in the coordination system
        3. Acquires slot leases for tag management
        4. Processes the acquired tags and establishes OPC subscriptions

        The preparation phase is critical for establishing the ingestor's role in
        the distributed system and ensuring it can begin processing OPC data.

        Raises:
            Exception: If any error occurs during the initialization or lease acquisition process.
                      This will cause the application to exit as the ingestor cannot function
                      without proper initialization.
        """

        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 active
        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 handles the case where an ingestor is active but has no
        allocated slots. It attempts to acquire a slot lease if slots are
        available in the system.

        Args:
            number_of_slots (int): The number of available slots in the system.

        Behavior:
            - Only attempts to acquire slots if the ingestor is currently active
              (has no managed tags) and there are slots available
            - Requests a single slot lease to begin processing
        """

        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.

        This method implements the load balancing logic for distributing OPC tag
        processing across multiple ingestor instances. It ensures optimal resource
        utilization and fair distribution of work.

        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 inactive, so there are "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 new managed tags and handles configuration changes.

        This method compares the current managed tags with the previous state and
        performs necessary operations to maintain synchronization:
        - Subscribes to newly allocated tags
        - Resubscribes to tags with changed configurations
        - Unsubscribes from deallocated tags

        Args:
            old_managed_tags (Dict[str, Any]): The previous state of managed tags.

        Behavior:
            - Compares current and previous tag configurations
            - Establishes subscriptions for new tags
            - Updates subscriptions for modified tags
            - Removes subscriptions for deallocated tags
            - Updates metrics to reflect current state
        """

        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 processing loop for managing ingestors and slots.

        This method is the core of the ingestor's operation, performing the following
        tasks in each iteration:
        1. Declares the ingestor as active to maintain its presence in the system
        2. Polls for slot updates and manages resource allocation
        3. Handles scenarios where no slots are available
        4. Manages slot leases based on system load and available resources
        5. Updates OPC server configurations and checks server integrity
        6. Synchronizes managed tags with the current system state

        The loop implements a sophisticated load balancing algorithm that:
        - Distributes OPC tag processing across multiple ingestor instances
        - Ensures optimal resource utilization
        - Maintains system stability during scaling operations
        - Provides real-time monitoring and metrics collection

        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)