src/utils/session_mapping.py

#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# 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