pip install loguru tenacity pandas matplotlib seaborn

import os
import signal
import time
import sys
import random
from collections import deque
import threading

from loguru import logger
from tenacity import retry, wait_exponential, stop_after_attempt, after_log
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def create_shutdown_state() -> dict:
    """
    Initializes a dictionary to manage the application's shutdown state.

    The state includes a flag to signal shutdown, a list of active tasks,
    and a record of cleanup operations.

    Parameters
    ----------
    None

    Returns
    -------
    dict
        An initialized dictionary with shutdown-related keys:
        - 'shutdown_flag': bool, indicates if shutdown has been requested.
        - 'active_tasks': list, tracks currently running tasks.
        - 'cleanup_log': list, records of cleanup operations.
        - 'start_time': float, timestamp when the state was created.
    """
    logger.info("Initializing shutdown state.")
    state = {
        'shutdown_flag': False,
        'active_tasks': [],
        'cleanup_log': [],
        'start_time': time.monotonic(),
        'shutdown_requested_time': None
    }
    logger.debug(f"Shutdown state initialized: {state}")
    return state

def _handle_signal(state: dict, signum: int, frame: object) -> None:
    """
    Internal signal handler to set the shutdown flag in the application state.

    Parameters
    ----------
    state : dict
        The current application state dictionary, which includes the 'shutdown_flag'.
    signum : int
        The number of the signal received (e.g., signal.SIGINT, signal.SIGTERM).
    frame : object
        The current stack frame (required by signal handler signature, but unused).

    Returns
    -------
    None

    Examples
    --------
    >>> app_state = create_shutdown_state()
    >>> # In a real scenario, this would be called by the OS when a signal is received.
    >>> _handle_signal(app_state, signal.SIGINT, None)
    >>> app_state['shutdown_flag']
    True
    """
    signal_name = signal.Signals(signum).name
    logger.warning(f"Received signal {signal_name} ({signum}). Initiating graceful shutdown.")
    state['shutdown_flag'] = True
    state['shutdown_requested_time'] = time.monotonic()
    logger.debug("Shutdown flag set to True.")

def register_signal_handlers(state: dict) -> dict:
    """
    Registers signal handlers for SIGINT (Ctrl+C) and SIGTERM (kill command).

    When these signals are received, the `_handle_signal` function will be called
    to set the shutdown flag in the provided state dictionary.

    Parameters
    ----------
    state : dict
        The application state dictionary where the shutdown flag will be set.

    Returns
    -------
    dict
        The updated state dictionary with signal handlers registered.

    Examples
    --------
    >>> app_state = create_shutdown_state()
    >>> updated_state = register_signal_handlers(app_state)
    >>> # Now, if SIGINT or SIGTERM are sent, updated_state['shutdown_flag'] will become True.
    """
    logger.info("Registering signal handlers for SIGINT and SIGTERM.")
    # Use a lambda to pass the state dictionary to the signal handler
    signal.signal(signal.SIGINT, lambda s, f: _handle_signal(state, s, f))
    signal.signal(signal.SIGTERM, lambda s, f: _handle_signal(state, s, f))
    logger.debug("Signal handlers registered successfully.")
    return state

@retry(
    wait=wait_exponential(multiplier=1, min=0.5, max=5), # Exponential backoff with min 0.5s, max 5s
    stop=stop_after_attempt(3), # Retry up to 3 times
    after=after_log(logger, 'warning') # Log retries as warnings
)
def perform_cleanup_task(state: dict, task_name: str, duration: float, fail_chance: float = 0.2) -> dict:
    """
    Simulates a cleanup task that might fail and be retried, with random jitter.

    Parameters
    ----------
    state : dict
        The current application state dictionary.
    task_name : str
        A descriptive name for the cleanup task.
    duration : float
        The simulated time (in seconds) the task takes.
    fail_chance : float, optional
        Probability (0.0 to 1.0) of the task failing once. Defaults to 0.2.

    Returns
    -------
    dict
        The updated application state dictionary, with the cleanup task logged.

    Raises
    ------
    RuntimeError
        If the cleanup task fails after all retries.

    Examples
    --------
    >>> app_state = create_shutdown_state()
    >>> # Example of a successful cleanup task
    >>> updated_state = perform_cleanup_task(app_state, "Close DB Connection", 1.0, fail_chance=0.0)
    >>> # Example of a potentially failing cleanup task (requires mocking or careful testing)
    >>> # updated_state = perform_cleanup_task(app_state, "Flush Logs", 0.5, fail_chance=0.8)
    """
    jitter = random.uniform(0, 0.1) # Add random jitter to simulate variability
    actual_duration = duration + jitter

    logger.info(f"Attempting cleanup task: '{task_name}' (simulated duration: {actual_duration:.2f}s).")

    if random.random() < fail_chance: # Simulate occasional failure
        logger.warning(f"Cleanup task '{task_name}' failed temporarily. Retrying...")
        raise RuntimeError(f"Simulated failure for task '{task_name}'")

    time.sleep(actual_duration)
    log_entry = {
        'task': task_name,
        'status': 'completed',
        'time_taken': actual_duration,
        'timestamp': time.monotonic() - state['start_time']
    }
    state['cleanup_log'].append(log_entry)
    logger.success(f"Cleanup task '{task_name}' completed successfully.")
    return state

def simulate_workload(state: dict, task_duration: float = 0.1) -> dict:
    """
    Simulates an ongoing application workload that checks for a shutdown signal.

    It adds tasks to `state['active_tasks']` until `state['shutdown_flag']` is set.

    Parameters
    ----------
    state : dict
        The current application state dictionary, containing 'shutdown_flag'
        and 'active_tasks'.
    task_duration : float, optional
        The simulated time (in seconds) each small unit of work takes. Defaults to 0.1.

    Returns
    -------
    dict
        The updated state dictionary after the workload has finished.

    Examples
    --------
    >>> app_state = create_shutdown_state()
    >>> # In a separate thread or process, you'd call this:
    >>> # t = threading.Thread(target=simulate_workload, args=(app_state, 0.05))
    >>> # t.start()
    >>> # time.sleep(1) # Let it run for a bit
    >>> # app_state['shutdown_flag'] = True # Signal shutdown
    >>> # t.join() # Wait for it to finish
    """
    task_count = 0
    logger.info("Starting simulated workload. Will run until shutdown is requested.")
    while not state['shutdown_flag']:
        jitter = random.uniform(0, 0.05) # Add small random jitter
        time.sleep(task_duration + jitter)
        task_count += 1
        state['active_tasks'].append(f"task_{task_count}")
        if task_count % 10 == 0:
            logger.debug(f"Workload active: processed {task_count} tasks.")

    logger.info(f"Simulated workload received shutdown signal. Processed {task_count} tasks.")
    return state

def run_graceful_shutdown(state: dict, workload_duration: float = 5, cleanup_timeout: float = 10) -> dict:
    """
    Orchestrates the graceful shutdown process for the application.

    This involves registering signal handlers, running a simulated workload,
    and then performing cleanup tasks once a shutdown is requested or triggered.

    Parameters
    ----------
    state : dict
        The application state dictionary.
    workload_duration : float, optional
        The simulated time (in seconds) the main workload runs before
        a shutdown is internally triggered (if no external signal is sent).
        Defaults to 5.
    cleanup_timeout : float, optional
        The maximum time (in seconds) allowed for all cleanup tasks to complete.
        Defaults to 10.

    Returns
    -------
    dict
        The final application state dictionary after the shutdown process.

    Examples
    --------
    >>> app_state = create_shutdown_state()
    >>> final_state = run_graceful_shutdown(app_state, workload_duration=3, cleanup_timeout=5)
    >>> # The 'cleanup_log' in final_state will show executed cleanup tasks.
    """
    logger.info("Starting graceful shutdown orchestration.")

    # 1. Register signal handlers
    state = register_signal_handlers(state)

    # 2. Start a simulated workload in a separate thread
    workload_thread = threading.Thread(target=simulate_workload, args=(state, 0.05))
    workload_thread.start()
    logger.info(f"Workload thread started. Will run for ~{workload_duration}s or until signal.")

    # 3. Wait for the shutdown flag or a simulated duration
    start_waiting = time.monotonic()
    while not state['shutdown_flag'] and (time.monotonic() - start_waiting < workload_duration):
        time.sleep(0.1) # Polling for shutdown flag

    if not state['shutdown_flag']:
        logger.info("Simulated workload duration elapsed. Internally triggering shutdown.")
        state['shutdown_flag'] = True
        state['shutdown_requested_time'] = time.monotonic()
    else:
        logger.info("Shutdown signal received or duration elapsed, proceeding to cleanup.")

    # Wait for workload to acknowledge shutdown flag and finish its current cycle
    logger.info("Waiting for workload thread to finish acknowledging shutdown flag...")
    workload_thread.join(timeout=2) # Give it a little time to finish its loop
    if workload_thread.is_alive():
        logger.warning("Workload thread did not finish cleanly within timeout. Forcing continue.")

    # 4. Perform cleanup tasks with a global timeout
    logger.info(f"Starting cleanup phase with a timeout of {cleanup_timeout} seconds.")
    cleanup_start_time = time.monotonic()

    cleanup_tasks = [
        ("Close Database Connection", 1.5, 0.3), # Name, Duration, Fail Chance
        ("Flush Application Logs", 0.8, 0.1),
        ("Save In-progress Work", 2.0, 0.05),
        ("Release Network Socket", 1.0, 0.2)
    ]

    for task_name, task_duration, fail_chance in cleanup_tasks:
        remaining_time = cleanup_timeout - (time.monotonic() - cleanup_start_time)
        if remaining_time <= 0:
            logger.error(f"Cleanup timeout reached. Aborting remaining cleanup tasks. Failed to complete: {task_name}")
            break
        try:
            # Simulate potential task delays or failures
            state = perform_cleanup_task(state, task_name, task_duration, fail_chance)
        except RuntimeError as e:
            logger.error(f"Cleanup task '{task_name}' failed critically after retries: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred during cleanup task '{task_name}': {e}")

    cleanup_end_time = time.monotonic()
    total_cleanup_time = cleanup_end_time - cleanup_start_time
    logger.info(f"Cleanup phase completed in {total_cleanup_time:.2f} seconds.")

    if state['shutdown_requested_time']:
        total_shutdown_duration = cleanup_end_time - state['shutdown_requested_time']
        logger.info(f"Total graceful shutdown duration from signal to completion: {total_shutdown_duration:.2f} seconds.")

    logger.info("Graceful shutdown orchestration finished.")
    return state

# Configure Loguru to write to stderr and capture it for display
logger.remove()
logger.add(sys.stderr, format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")

logger.info("### Starting Graceful Shutdown Demonstration ###")

# Initialize the application state
app_state = create_shutdown_state()

# Run the graceful shutdown process
# You can try interrupting this cell (Ctrl+C) while it's running to see the signal handler in action.
final_state = run_graceful_shutdown(app_state, workload_duration=3, cleanup_timeout=7)

logger.info("### Graceful Shutdown Demonstration Complete ###")

# Display summary of cleanup activities
print("\n--- Cleanup Summary ---")
if final_state['cleanup_log']:
    cleanup_df = pd.DataFrame(final_state['cleanup_log'])
    display(cleanup_df)
else:
    print("No cleanup tasks were logged.")

print(f"\nTotal active tasks before shutdown: {len(app_state['active_tasks'])}")
print(f"Shutdown flag final state: {final_state['shutdown_flag']}")

if final_state['cleanup_log']:
    cleanup_df = pd.DataFrame(final_state['cleanup_log'])

    # Calculate time relative to shutdown request for visualization
    # If shutdown was internal, use state creation time as reference
    reference_time = final_state['shutdown_requested_time'] if final_state['shutdown_requested_time'] else final_state['start_time']
    cleanup_df['time_relative_to_shutdown'] = cleanup_df['timestamp'] - (reference_time - final_state['start_time'])

    plt.figure(figsize=(10, 6))
    sns.barplot(x='task', y='time_taken', data=cleanup_df, palette='viridis')
    plt.axvline(x=-0.5, color='red', linestyle='--', label='Shutdown Signal/Trigger')
    plt.title('Cleanup Task Durations During Graceful Shutdown')
    plt.xlabel('Cleanup Task')
    plt.ylabel('Time Taken (seconds)')
    plt.xticks(rotation=45, ha='right')
    plt.legend()
    plt.tight_layout()
    plt.show()

    # Visualize cleanup progress over time
    cleanup_df = cleanup_df.sort_values('time_relative_to_shutdown')
    cleanup_df['cumulative_time'] = cleanup_df['time_taken'].cumsum()

    plt.figure(figsize=(12, 6))
    plt.plot(cleanup_df['time_relative_to_shutdown'], cleanup_df['cumulative_time'], marker='o', linestyle='-', color='skyblue')
    plt.axvline(x=0, color='red', linestyle='--', label='Shutdown Initiated')
    plt.axhline(y=7, color='green', linestyle=':', label='Cleanup Timeout (7s)')
    plt.title('Cumulative Cleanup Time After Shutdown Initiation')
    plt.xlabel('Time Relative to Shutdown Initiation (seconds)')
    plt.ylabel('Cumulative Time Spent on Cleanup (seconds)')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend()
    plt.tight_layout()
    plt.show()
else:
    print("No data to visualize cleanup activities.")