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

import asyncio
import traceback
from copy import deepcopy

from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring

import ingestor.metrics as metrics
from ingestor.managers.data_manager import DataManager
from ingestor.managers.opc_manager import OpcManager
from ingestor.managers.resource_manager import ResourceManager


class IngestorManager(SientiaMonitoring):
    """
    Central coordinator for managing OPC data ingestion operations.

    The IngestorManager orchestrates the interaction between different components:
    - DataManager: Handles data persistence and Kafka export
    - OPC Managers: Manage individual OPC UA server connections
    - ResourceManager: Coordinates slot leasing and load balancing

    This class implements a slot-based architecture where:
    - Each slot represents a collection of OPC tags from one or more servers
    - Slots are distributed across multiple ingestor instances for load balancing
    - Dynamic slot allocation ensures optimal resource utilization

    Key Responsibilities:
    - Slot lease management and distribution
    - OPC server connection lifecycle management
    - Tag subscription coordination
    - System health monitoring and integrity checks
    - Load balancing across multiple ingestor instances

    Args:
        kafka_servers (str): Comma-separated list of Kafka server addresses
        redis_data (dict): Redis connection parameters (host, port, username, password)
        lease_ttl (int): Time-to-live for slot leases in seconds
        heartbeat_ttl (int): Time-to-live for heartbeat signals in seconds
        poll_interval (int): Main loop polling interval in seconds
        mongo_connection_string (str): MongoDB connection string
        mongo_database (str): MongoDB database name
        metadata (dict): Application metadata for notifications and tracking
        logger (Logger): Logger instance for application logging
        notification_handler (NotificationHandler): Handler for sending notifications
        export_to_kafka (bool): Whether to export data to Kafka

    Attributes:
        data_manager (DataManager): Manages data persistence and Kafka export
        opc_managers (dict): Dictionary of OPC managers keyed by server name
        resource_manager (ResourceManager): Manages Redis-based resource coordination
        number_of_slots (int): Total number of slots configured in the system
        poll_interval (int): Main loop polling interval
        managed_tags (dict): Currently managed tags organized by slot
        opc_servers (dict): OPC server configurations
        metadata (dict): Application metadata
    """

    def __init__(
        self,
        kafka_servers: str,
        redis_data: dict,
        lease_ttl: int,
        heartbeat_ttl: int,
        poll_interval: int,
        mongo_connection_string: str,
        mongo_database: str,
        metadata: dict,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
        export_to_kafka: bool = False,
    ):
        redis_host: str = redis_data['host']
        redis_port: int = int(redis_data['port'])
        redis_username: str | None = redis_data.get('username', None)
        redis_password: str | None = redis_data.get('password', None)

        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        self.data_manager = DataManager(
            kafka_servers=kafka_servers,
            mongo_connection_string=mongo_connection_string,
            mongo_database=mongo_database,
            export_to_kafka=export_to_kafka,
            metadata=metadata,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )
        self.opc_managers: dict = {}
        self.resource_manager = ResourceManager(
            host=redis_host,
            port=redis_port,
            lease_ttl=lease_ttl,
            heartbeat_ttl=heartbeat_ttl,
            metadata=metadata,
            logger=logger,
            notification_handler=notification_handler,
            username=redis_username,
            password=redis_password,
            metrics_controller=metrics_controller,
        )
        self.number_of_slots = 0
        self.poll_interval = poll_interval
        self.managed_tags: dict = {}
        self.opc_servers: dict = {}

        self.metadata = metadata

    async def initialize_opc_from_config(self, server_config: dict) -> OpcManager | None:
        """
        Initializes an OPC Manager instance using the provided server configuration.

        This method creates and configures an OPC Manager for a specific OPC UA server,
        establishing the connection and preparing it for tag subscriptions.

        Args:
            server_config (dict): A dictionary containing the OPC server configuration.
                Expected keys include:
                    - 'name' (str): The name of the OPC server.
                    - 'url' (str): The URL of the OPC server.
                    - 'server_uri' (str): The URI of the OPC server.
                    - 'cert_path' (str, optional): Path to the client certificate file.
                    - 'private_key_path' (str, optional): Path to the private key file.
                    - 'server_cert_path' (str, optional): Path to the server certificate file.

        Returns:
            OpcManager | None: An initialized OpcManager instance if successful,
            otherwise None if an error occurs during initialization.

        Raises:
            Exception: If OPC manager initialization fails, the error is logged and
                      a notification is sent, but the method returns None to allow
                      the system to continue operating with other servers.
        """

        try:
            self.logger.info(f'Initializing OpcManager at {server_config["url"]}')
            manager = OpcManager(
                name=server_config['name'],
                url=server_config['url'],
                subscription_period_ms=server_config['subscription_period_ms'],
                data_manager=self.data_manager,
                logger=self.logger,
                server_uri=server_config['server_uri'],
                notification_handler=self.notification_handler,
                metadata=self.metadata,
                cert_path=server_config.get('cert_path'),
                private_key_path=server_config.get('private_key_path'),
                server_cert_path=server_config.get('server_cert_path'),
                metrics_controller=self.metrics_controller,
            )

            manager.config = server_config
            await manager.connect()
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=self.metadata,
                notification_id=f'OPC_CONNECTION_ERROR_{server_config["name"]}',
                message=f'Error initializing OPC manager: {e}',
                block='opc_manager',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            self.logger.error(trace)

            return None

        return manager

    async def shutdown(self):
        """
        Gracefully shuts down the IngestorManager and all its components.

        This method ensures proper cleanup of:
        - All OPC manager instances and their connections
        - Data manager connections and resources
        - Active subscriptions and server connections

        The shutdown process is performed asynchronously to allow proper cleanup
        of all managed resources before termination.
        """
        for _server_name, server in self.opc_managers.items():
            await server.shutdown()
        self.data_manager.shutdown()

    def __del__(self):
        asyncio.run(self.shutdown())

    async def remove_server(self, server: str):
        """
        Removes an OPC server from the ingestor.
        """
        if server in self.opc_managers:
            await self.opc_managers[server].shutdown()
            del self.opc_managers[server]
        for slot, _config in self.managed_tags.items():
            self.managed_tags[slot].pop(server, None)

    async def update_opc_servers(self):
        """
        Updates the OPC (OLE for Process Control) server connections managed by the ingestor.
        This method ensures that the OPC servers defined in `self.managed_tags` are properly
        initialized and updated. It performs the following tasks:
        - Registers new OPC servers based on the configuration in `self.managed_tags`.
        - Updates existing OPC server instances if their configuration has changed.
        - Disconnects and removes OPC servers that are no longer present in `self.managed_tags`.
        Steps:
        1. Iterates through the `self.managed_tags` dictionary to identify and register servers.
        2. Initializes new OPC server instances if they are not already managed.
        3. Reinitializes OPC server instances if their configuration has changed.
        4. Disconnects and removes OPC servers that are no longer registered.
        Attributes:
            self.managed_tags (dict): A nested dictionary containing slot and server configurations.
            self.opc_managers (dict): A dictionary mapping server names to their
                OPC manager instances.
            self.data_manager: An object responsible for managing data operations.
            self.logger: A logging object for recording warnings and other messages.
        Raises:
            Any exceptions raised during OPC server initialization or disconnection.
        Logs:
            - Warnings for servers that are no longer found in `self.managed_tags`.
        """

        registered_servers = []
        current_managed_tags = deepcopy(self.managed_tags)
        for _slot, slot_config in current_managed_tags.items():
            for server, server_config in slot_config.items():
                registered_servers.append(server)
                server_config = deepcopy(server_config)
                server_config.pop('tags', None)
                server_instance = self.opc_managers.get(server, None)
                if server_instance is None:
                    self.logger.info(f'Initializing OPC manager for server {server}')
                    server_instance = await self.initialize_opc_from_config(server_config)

                elif server_instance.config != server_config:
                    self.logger.warning(f'Reinitializing OPC manager for server {server}')
                    await server_instance.shutdown()
                    del self.opc_managers[server]
                    server_instance = await self.initialize_opc_from_config(server_config)
                else:
                    self.logger.debug(
                        f'OPC manager for server {server} is already initialized and up to date'
                    )

                if server_instance is not None:
                    self.opc_managers[server] = server_instance
                else:
                    self.logger.warning(
                        f'Failed to initialize OPC manager for server {server}, '
                        f'removing server from managed tags.'
                    )

                    await self.remove_server(server)

        servers = list(self.opc_managers.keys())
        for server in servers:
            if server not in registered_servers:
                self.logger.warning(
                    f'Server {server} not found in managed tags. Desconnecting from server.'
                )
                await self.remove_server(server)

        await self.emit_metric(
            metric_object=metrics.OPC_MANAGERS_ACTIVE,
            method='set',
            value=len(self.opc_managers),
            tags={
                'pod_id': self.pod_id,
            },
        )

    async def check_opc_servers_integrity(self):
        """
        Checks the integrity of the OPC servers and updates the OPC servers if necessary.

        This method performs health checks on all managed OPC servers by:
        - Checking cycle counts for data reception
        - Monitoring connection health and data flow
        - Triggering reconnection for lost servers
        - Updating metrics for active OPC managers

        Side Effects:
            - Updates cycle monitoring for all nodes
            - Removes lost servers from managed tags
            - Updates OPC manager metrics
        """
        to_disconnect: list[str] = []
        for server, opc_manager in self.opc_managers.items():
            opc_manager.check_cycles()

            is_lost = opc_manager.check_opc_listenning()
            if is_lost:
                self.logger.warning(f'OPC server {server} is lost. Server will be disconnected.')

                to_disconnect.append(server)

        for server in to_disconnect:
            await self.remove_server(server)

        await self.emit_metric(
            metric_object=metrics.OPC_MANAGERS_ACTIVE,
            method='set',
            value=len(self.opc_managers),
            tags={
                'pod_id': self.pod_id,
            },
        )

    async def declare_active(self):
        """
        Declares the ingestor as active by sending a heartbeat signal to the resource manager.

        This method ensures that the ingestor is marked as active by invoking the
        `ingestor_heartbeat` method of the associated resource manager. The heartbeat
        mechanism enables load balancers and monitoring systems to track active instances.

        Side Effects:
            - Updates Redis with current instance heartbeat
            - Enables load balancing and health monitoring
        """

        await self.resource_manager.ingestor_heartbeat()

    async def get_active_ingestors(self) -> list[str]:
        """
        Retrieve a list of active ingestors.

        This method fetches all ingestors from the resource manager and returns them.
        If no ingestors are found, an empty list is returned.

        Returns:
            List[str]: A list of active ingestor names, or an empty list if none are found.

        The method queries Redis for all active ingestor heartbeats and extracts
        the pod identifiers for load balancing and coordination purposes.
        """

        ingestors = await self.resource_manager.get_all_ingestors()
        return ingestors if ingestors else []

    async def get_number_of_leases(self) -> int:
        """
        Retrieves the number of leases managed by the resource manager.

        This method fetches all available leases from the resource manager,
        calculates their count, and updates the `number_of_slots` attribute.

        Returns:
            int: The total number of leases. Returns 0 if no leases are available.

        Side Effects:
            - Updates internal slot count tracking
            - Updates Prometheus metrics for total leases
        """

        leases = await self.resource_manager.get_all_leases()
        self.number_of_slots = len(leases) if leases else 0
        await self.emit_metric(
            metric_object=metrics.LEASES_TOTAL,
            method='set',
            value=self.number_of_slots,
            tags={
                'pod_id': self.pod_id,
            },
        )
        return self.number_of_slots

    async def get_number_of_slots(self) -> int:
        """
        Retrieves the number of slots managed by the resource manager.

        This method fetches all available slots from the resource manager,
        calculates their count, and updates the `number_of_slots` attribute.

        Returns:
            int: The total number of slots. Returns 0 if no slots are available.

        Side Effects:
            - Updates internal slot count tracking
            - Updates Prometheus metrics for total slots
        """

        slots = await self.resource_manager.get_all_slots()
        self.number_of_slots = len(slots) if slots else 0
        await self.emit_metric(
            metric_object=metrics.SLOTS_TOTAL,
            method='set',
            value=self.number_of_slots,
            tags={
                'pod_id': self.pod_id,
            },
        )
        return self.number_of_slots

    async def get_slot_leases(self, max_slots: int = 1) -> dict:
        """
        Acquires a specified number of resource slots by leasing them from the resource manager.

        This method implements the slot acquisition logic for load balancing:
        - Iterates through available slots and attempts to lease them
        - Logs the leasing of each slot
        - Updates the `managed_tags` attribute with the acquired slots
        - Stops leasing once the specified `max_slots` are acquired

        Args:
            max_slots (int): The maximum number of slots to lease. Defaults to 1.

        Returns:
            Dict: A dictionary where the keys are the slot identifiers (as strings)
                  and the values are the leased slot details.

        Behavior:
            - Attempts to lease slots sequentially starting from slot 1
            - Skips slots that cannot be retrieved after leasing
            - Logs warnings if unable to acquire the requested number of slots
            - Updates metrics for acquired slots and managed slots count

        Notes:
            - If a slot is leased but its details cannot be retrieved
              (i.e., `get_tag_slot` returns None), that slot is skipped.
        """

        acquired = {}
        for i in range(1, self.number_of_slots + 1):
            if await self.resource_manager.lease_tag(str(i)):
                self.logger.info(f'Leased slot {i}')
                slots = await self.resource_manager.get_tag_slot(str(i))
                if slots is None:
                    continue
                acquired[str(i)] = slots
                await self.emit_metric(
                    metric_object=metrics.SLOTS_ACQUIRED,
                    method='inc',
                    value=1,
                    tags={
                        'pod_id': self.pod_id,
                    },
                )

                if len(acquired) >= max_slots:
                    self.managed_tags.update(acquired)
                    await self.emit_metric(
                        metric_object=metrics.SLOTS_MANAGED,
                        method='set',
                        value=len(self.managed_tags),
                        tags={
                            'pod_id': self.pod_id,
                        },
                    )
                    return acquired

        self.logger.warning(
            f'Unable to acquire {max_slots} slots. Only {acquired} slots were leased.'
        )
        self.managed_tags.update(acquired)
        await self.emit_metric(
            metric_object=metrics.SLOTS_MANAGED,
            method='set',
            value=len(self.managed_tags),
            tags={
                'pod_id': self.pod_id,
            },
        )
        return acquired

    async def unsubscribe_slot(self, slot: str):
        """
        Unsubscribes a specific slot from all associated OPC servers.

        This method removes all subscriptions for a given slot across all
        OPC servers that were managing it. It ensures clean cleanup of
        resources when slots are released or reconfigured.

        Args:
            slot (str): The name of the slot to unsubscribe.

        Raises:
            KeyError: If the specified slot does not exist in the managed tags.

        Side Effects:
            - Removes subscriptions from all OPC servers for the specified slot
            - Cleans up subscription resources on the OPC servers
        """

        for server in self.managed_tags[slot].keys():
            if server in self.opc_managers:
                await self.opc_managers[server].unsubscribe(slot)

    async def update_slot_config(self):
        """
        Updates the configuration of managed slots by renewing their leases,
        fetching the latest configurations, and handling any changes or removals.

        This method performs the following steps:
        1. Renews the lease for each managed slot using the resource manager.
        2. Fetches the latest configuration for each slot.
        3. Logs and removes slots whose configurations are no longer available.
        4. Updates the configuration of slots if changes are detected.
        5. Unsubscribes and re-subscribes to slots with updated configurations.
        6. Removes slots from the managed tags if they are no longer valid.
        7. Updates the OPC servers after processing all slots.

        Side Effects:
            - Modifies the `managed_tags` dictionary to reflect the latest slot configurations.
            - Updates OPC server subscriptions based on the current state of managed slots.

        Raises:
            None explicitly, but relies on the behavior of `resource_manager` and
            other dependencies for error handling.

        Logging:
            - Logs warnings for removed slots.
            - Logs informational messages for updated slot configurations.
        """

        removed_slots: list[str] = []
        for slot, _slot_config in self.managed_tags.items():
            await self.resource_manager.renew_tag_lease(slot)
            update = await self.resource_manager.get_tag_slot(slot)
            if update is None:
                removed_slots.append(slot)
                continue

            self.managed_tags[slot] = update

        for slot in removed_slots:
            self.managed_tags.pop(slot, None)

    async def drop_slot_leases(self, ids: list[str]) -> None:
        """
        Releases the leases associated with the specified slot IDs.

        This method iterates through a list of slot IDs and calls the
        `drop_tag_lease` method of the `resource_manager` to release
        the lease for each ID. It's used during load balancing and
        graceful shutdown scenarios.

        Args:
            ids (List[str]): A list of slot IDs for which the leases
                             should be released.

        Returns:
            None

        Side Effects:
            - Releases Redis-based leases for specified slots
            - Updates metrics for released slots count
        """
        for lease_id in ids:
            await self.resource_manager.drop_tag_lease(lease_id)
            await self.emit_metric(
                metric_object=metrics.SLOTS_RELEASED,
                method='inc',
                value=1,
                tags={
                    'pod_id': self.pod_id,
                },
            )

    async def manage_server(self, slot: str, server: str, server_config: dict, tags: dict) -> int:
        """
        Manages the subscription of tags to a specified OPC server and slot.

        This method ensures that the specified server and slot have an active subscription
        for the provided tags. If the server or slot is not properly configured, or if
        subscription fails, appropriate error handling is performed.

        Args:
            slot (str): The slot identifier for the subscription.
            server (str): The name of the OPC server.
            server_config (dict): Configuration dictionary for the server, which includes
                the tags to be subscribed under the key 'tags'.
            tags (dict): A dictionary of tags to be subscribed.

        Returns:
            int: Status code indicating the result of the operation:
                - 0: Subscription was successful.
                - 1: Server not found in `opc_managers`.
                - 2: Subscription creation or tag subscription failed.

        Logs:
            - Logs informational messages about the subscription process.
            - Logs errors if the server is not found, subscription creation fails, or
              tag subscription fails.
            - Logs a warning if a subscription is removed due to failure.

        Raises:
            Exception: Any unexpected exceptions during subscription creation or tag
            subscription are logged but not propagated.

        Side Effects:
            - Creates or updates OPC subscriptions
            - Manages tag subscriptions on OPC servers
            - Updates error metrics and notifications
        """

        self.logger.info(f'Subscribing to tags from {slot}:{server}')
        tags_to_sub = server_config.get('tags')
        if server not in self.opc_managers:
            self.logger.error(f'Server {server} not found in opc_managers.')
            return 1
        if slot not in self.opc_managers[server].subscriptions:
            try:
                await self.opc_managers[server].create_subscription(slot)
            except Exception as e:
                self.logger.error(f'Failed to create subscription for slot {slot}: {e}')
                return 2
        try:
            self.logger.info(tags_to_sub)
            await self.opc_managers[server].subscribe(
                slot, deepcopy(tags_to_sub), self.poll_interval
            )
            self.logger.info(tags_to_sub)
        except Exception as e:
            await self.emit_metric(
                metric_object=metrics.OPC_SUBSCRIPTION_ERRORS,
                tags={
                    'pod_id': self.pod_id,
                    'server': server,
                    'slot': slot,
                },
            )
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=self.metadata,
                notification_id=f'OPC_SUBSCRIPTION_ERROR_{slot}:{server}',
                message=f'Failed to subscribe to tags from {slot}:{server}\n{tags}: {e}',
                block='opc_manager',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            self.logger.error(trace)

            self.logger.warning(f'Removing subscription from server {server} for slot {slot}')
            await self.opc_managers[server].unsubscribe(slot)
            return 2
        return 0

    async def subscribe_to_tags(self, tags: dict) -> None:
        """
        Subscribes to a set of tags and manages their configurations.

        This method processes a dictionary of tags, iterating through each slot and server
        configuration. It attempts to manage the server configurations and removes any
        servers that return a specific response code.

        Args:
            tags (Dict): A dictionary containing tag configurations. The structure is
                         expected to be {slot: {server: server_config}}.

        Side Effects:
            - Logs the provided tags for debugging purposes.
            - Updates the `managed_tags` attribute by removing servers that meet the
              removal criteria.
            - Establishes OPC subscriptions for all configured tags.

        Removal Criteria:
            - If the `manage_server` method returns a response code of 2 for a given
              slot and server, that server is removed from the `managed_tags` attribute.

        The method ensures that only successfully configured servers remain in the
        managed tags, maintaining system stability and preventing subscription errors.
        """

        to_remove = []
        self.logger.info(tags)
        for slot, slot_config in tags.items():
            for server, server_config in slot_config.items():
                response = await self.manage_server(slot, server, server_config, tags)
                if response == 2:
                    to_remove.append([slot, server])

        for slot, server in to_remove:
            if server in self.opc_managers:
                await self.opc_managers[server].shutdown()
                del self.opc_managers[server]
            self.managed_tags[slot].pop(server, None)