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() def __del__(self): asyncio.run(self.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)