# # SPDX-FileCopyrightText: Hadad # SPDX-License-Identifier: Apache-2.0 # import random # Import random module to enable random selection from a list from datetime import datetime # Import datetime class to work with current UTC time from typing import Dict, List # Import type hints for dictionaries and lists (not explicitly used here but imported) from config import auth # Import authentication configuration, likely a list of host dictionaries with credentials from src.utils.helper import busy, mark # Import 'busy' dictionary and 'mark' function to track and update host busy status # Initialize a global dictionary to map session IDs to assigned hosts mapping = {} # Store session_id to host assignment mapping to maintain consistent host allocation per session # Define a function to get an available host for a given session, optionally excluding certain hosts def get_host(session_id: str, exclude_hosts: List[str] = None) -> dict: """ Retrieve or assign a host for the given session ID. Args: session_id (str): A unique identifier for the current session. Returns: dict: The selected host dictionary from the auth configuration. Raises: Exception: If no available hosts are found to assign. Explanation: Retrieve an available host for the specified session ID, ensuring excluded hosts are not assigned. This function maintains a mapping of session IDs to hosts to provide consistent host assignment. It filters out busy hosts and those explicitly excluded, then randomly selects an available host. If no hosts are available, it raises an exception. """ # If no list of hosts to exclude is provided, initialize it as an empty list if exclude_hosts is None: # Check if exclude_hosts parameter was omitted or set to None exclude_hosts = [] # Initialize exclude_hosts to an empty list to avoid errors during filtering # Check if the session ID already has an assigned host in the mapping dictionary if session_id in mapping: # Verify if a host was previously assigned to this session assigned_host = mapping[session_id] # Retrieve the assigned host dictionary for this session # If the assigned host is not in the list of hosts to exclude, return it immediately if assigned_host["jarvis"] not in exclude_hosts: # Ensure assigned host is allowed for this request return assigned_host # Return the cached host assignment for session consistency else: # If the assigned host is excluded, remove the mapping to allow reassignment del mapping[session_id] # Delete the existing session-host mapping to find a new host # Get the current UTC time to compare against host busy status timestamps now = datetime.utcnow() # Capture current time to filter out hosts that are still busy # Create a list of hosts that are not currently busy and not in the exclude list available_hosts = [ h for h in auth # Iterate over all hosts defined in the authentication configuration if h["jarvis"] not in busy or busy[h["jarvis"]] <= now # Include hosts not busy or whose busy time has expired if h["jarvis"] not in exclude_hosts # Exclude hosts specified in the exclude_hosts list ] # If no hosts are available after filtering, raise an exception to indicate resource exhaustion if not available_hosts: # Check if the filtered list of hosts is empty raise Exception("No available hosts to assign.") # Inform caller that no hosts can be assigned currently # Randomly select one host from the list of available hosts to distribute load evenly selected = random.choice(available_hosts) # Choose a host at random to avoid bias in host selection # Store the selected host in the mapping dictionary for future requests with the same session ID mapping[session_id] = selected # Cache the selected host to maintain session affinity # Mark the selected host as busy using the helper function to update its busy status mark(selected["jarvis"]) # Update the busy dictionary to indicate this host is now in use # Return the selected host dictionary to the caller for use in processing the session return selected # Provide the caller with the assigned host details