api_server.py

from flask import Flask, request, jsonify
import threading
import time
import docker
import logging
import requests
from datetime import datetime
import socket
import random

app = Flask(__name__)

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("API_SERVER")

# Docker client
try:
    docker_client = docker.from_env()
    logger.debug("Connected to Docker")
except Exception as e:
    logger.warning(f"Docker unavailable: {e}")
    docker_client = None

# Cluster state
nodes = {}  # node_id -> dict
pods = {}   # pod_id -> dict

HEARTBEAT_TIMEOUT = 15  # seconds

# -----------------------------
# Utility Functions
# -----------------------------


def get_free_port():
    """Find a random available port on localhost."""
    while True:
        port = random.randint(8000, 9000)
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(("localhost", port))
                return port
            except OSError:
                continue  # Port is taken, try another

def create_node(node_id, cpu_cores):
    """Create a node container, find a free port, and register the node."""
    port = get_free_port()
    container_id = None

    if docker_client:
        try:
            container = docker_client.containers.run(
                'node-simulator',
                detach=True,
                environment={
                    'NODE_ID': node_id,
                    'API_URL': 'http://host.docker.internal:5000'
                },
                ports={'80/tcp': port},
                name=f"node-{node_id}",
                remove=True
            )
            container_id = container.id
            logger.info(f"Started container {container_id} for node {node_id}")
        except Exception as e:
            logger.warning(f"Container launch failed for {node_id}: {e}")

    # Register node
    nodes[node_id] = {
        "total_cpu": cpu_cores,
        "available_cpu": cpu_cores,
        "status": "UP",
        "last_heartbeat": time.time(),
        "container_id": container_id,
        "pods": [],
        "port": port,
        "created_at": datetime.now().isoformat()
    }

    return nodes[node_id]

def schedule_pod(pod_id, required_cpu):
    """Try to assign pod to the best node. Return True if success, False if pending/fail."""
    best_node = None
    best_score = None  # smallest leftover CPU

    for n_id, node in nodes.items():
        if node['status'] != 'UP':
            continue
        if node['available_cpu'] >= required_cpu:
            leftover = node['available_cpu'] - required_cpu
            if best_score is None or leftover < best_score:
                best_score = leftover
                best_node = n_id

    if best_node:
        # Assign pod
        nodes[best_node]['available_cpu'] -= required_cpu
        nodes[best_node]['pods'].append(pod_id)

        pods[pod_id].update({
            "node_id": best_node,
            "status": "🟢"
        })

        logger.info(f"Pod {pod_id} scheduled to node {best_node} (leftover CPU: {best_score})")

        # Notify node
        try:
            requests.post(
                f"http://localhost:{nodes[best_node]['port']}/pod",
                json={"pod_id": pod_id, "cpu": required_cpu},
                timeout=2
            )
        except Exception as e:
            logger.warning(f"Failed to notify node {best_node} for pod {pod_id}: {e}")

        return True

    # No suitable node
    pods[pod_id].update({
        "node_id": None,
        "status": "🟡" if nodes else "🔴"
    })

    logger.warning(f"Pod {pod_id} could not be scheduled. Status: {pods[pod_id]['status']}")
    return False

def reschedule_all_pods_from_node(node_id):
    """Reassign all pods from a failed node."""
    failed_pods = nodes[node_id]['pods'][:]
    nodes[node_id]['pods'].clear()
    nodes[node_id]['available_cpu'] = nodes[node_id]['total_cpu']
    for pod_id in failed_pods:
        schedule_pod(pod_id, pods[pod_id]['cpu'])

def retry_pending_pods():
    for pod_id, pod in pods.items():
        if pod['status'] == '🟡':
            logger.info(f"Retrying pending pod {pod_id}")
            schedule_pod(pod_id, pod['cpu'])

def remove_pod_from_node(node_id, pod_id):
    """Remove pod from a node and notify the simulator."""
    if pod_id in nodes[node_id]['pods']:
        nodes[node_id]['available_cpu'] += pods[pod_id]['cpu']
        nodes[node_id]['pods'].remove(pod_id)
        try:
            requests.delete(
                f"http://localhost:{nodes[node_id]['port']}/pod",
                json={"pod_id": pod_id},
                timeout=2
            )
        except Exception as e:
            logger.warning(f"Failed to deassign pod {pod_id} from node {node_id}: {e}")

# -----------------------------
# Heartbeat Thread
# -----------------------------

def health_checker():
    while True:
        current_time = time.time()

        for n_id, node in list(nodes.items()):
            # Node failure check
            if current_time - node['last_heartbeat'] > HEARTBEAT_TIMEOUT:
                if node['status'] != 'DOWN':
                    logger.warning(f"Node {n_id} marked DOWN (missed heartbeat)")
                    node['status'] = 'DOWN'
                    reschedule_all_pods_from_node(n_id)

            # Pending pod rescheduling retry
            for pod_id, pod in list(pods.items()):
                if pod['status'] == '🟡':
                    logger.info(f"Retrying pending pod {pod_id}")
                    schedule_pod(pod_id, pod['cpu'])

        time.sleep(5)

# -----------------------------
# Routes
# -----------------------------

@app.route('/add_node', methods=['POST'])
def add_node():
    data = request.get_json()
    n_id = data.get('node_id')
    cpu = data.get('cpu_cores')

    if not n_id or cpu is None:
        return jsonify({'error': 'Missing node_id or cpu_cores'}), 400
    if n_id in nodes:
        return jsonify({'error': 'Node already exists'}), 400

    try:
        cpu = int(cpu)
    except ValueError:
        return jsonify({'error': 'CPU must be an integer'}), 400
    if cpu <= 0:
        return jsonify({'error': 'CPU must be > 0'}), 400

    node_info = create_node(n_id, cpu)

    retry_pending_pods()

    return jsonify({
        'message': f"Node {n_id} added",
        'node': node_info
    }), 201

@app.route('/heartbeat', methods=['POST'])
def heartbeat():
    data = request.get_json()
    n_id = data.get("n_id")
    failed_pods = data.get("failed_pods", [])

    if not n_id or n_id not in nodes:
        return jsonify({'error': 'Invalid or unknown node_id'}), 400

    node = nodes[n_id]
    node['last_heartbeat'] = time.time()

    if node['status'] == 'DOWN':
        node['status'] = 'UP'
        logger.info(f"Node {n_id} recovered from DOWN")

    # Handle failed pods
    for pod_id in failed_pods:
        if pod_id in pods and pods[pod_id]['node_id'] == n_id:
            logger.warning(f"Pod {pod_id} failed on node {n_id}, rescheduling...")
            remove_pod_from_node(n_id, pod_id)
            schedule_pod(pod_id, pods[pod_id]['cpu'])

    return jsonify({'message': 'Heartbeat acknowledged'}), 200


@app.route('/launch_pod', methods=['POST'])
def launch_pod():
    data = request.get_json()
    pod_id = data.get('pod_id')
    cpu = data.get('cpu')

    if not pod_id or cpu is None:
        return jsonify({'error': 'Missing pod_id or cpu'}), 400
    if pod_id in pods:
        return jsonify({'error': 'Pod already exists'}), 400

    try:
        cpu = int(cpu)
    except ValueError:
        return jsonify({'error': 'CPU must be an integer'}), 400
    if cpu <= 0:
        return jsonify({'error': 'CPU must be > 0'}), 400

    pods[pod_id] = {
        "cpu": cpu,
        "status": "🟡",
        "node_id": None,
        "created_at": datetime.now().isoformat()
    }

    success = schedule_pod(pod_id, cpu)
    return jsonify({
        "message": "Pod launched",
        "status": pods[pod_id]['status'],
        "pod": pods[pod_id]
    }), 201 if success else 202


@app.route('/list_nodes', methods=['GET'])
def list_nodes():
    return jsonify([
        {
            "node_id": n_id,
            "status": node["status"],
            "cpu_total": node["total_cpu"],
            "cpu_available": node["available_cpu"],
            "pods": node["pods"],
            "last_heartbeat": node["last_heartbeat"]
        }
        for n_id, node in nodes.items()
    ])


@app.route('/list_pods', methods=['GET'])
def list_pods():
    return jsonify([
        {
            "pod_id": p_id,
            "cpu": pod["cpu"],
            "status": pod["status"],
            "node_id": pod.get("node_id"),
            "created_at": pod["created_at"]
        }
        for p_id, pod in pods.items()
    ])


# -----------------------------
# Server Init
# -----------------------------

if __name__ == "__main__":
    threading.Thread(target=health_checker, daemon=True).start()
    logger.info("API server started on port 5000")
    app.run(host="0.0.0.0", port=5000, debug=True)