api_server.py

from flask import Flask, request, jsonify
import docker
import json
from datetime import datetime, timedelta
import threading
import os
from pytz import UTC

app = Flask(__name__)
client = docker.from_env()
nodes = {}
track_pods = {}
policies = []  # For network policy simulation
pod_counter = 0
node_counter = 0
HEARTBEAT_THRESHOLD = 30  # Increased for testing
SCALE_UP_THRESHOLD = 0.2  # 20% available cores threshold
SCALE_DOWN_THRESHOLD = 0.5  # 50% utilization threshold
DEFAULT_NODE_CORES = 4  # Default cores for new nodes

# Load nodes and pods from JSON file on startup
def load_nodes():
    global nodes
    if os.path.exists("nodes.json"):
        with open("nodes.json", "r") as f:
            nodes = json.load(f)

# Save nodes and pods to JSON file
def save_nodes():
    with open("nodes.json", "w") as f:
        json.dump(nodes, f)

# Generate unique node ID
def generate_node_id():
    global node_counter
    node_counter += 1
    return f"node{node_counter}"

# Generate unique pod ID
def generate_pod_id():
    global pod_counter
    pod_counter += 1
    return f"pod{pod_counter}"

# Compute node status
def get_node_status(last_heartbeat):
    last_time = datetime.fromisoformat(last_heartbeat)
    current_time = datetime.now(UTC)
    return "healthy" if (current_time - last_time) < timedelta(seconds=HEARTBEAT_THRESHOLD) else "failed"

# Check node health and handle failures
def check_node_health():
    global track_pods 
    failed_nodes = []
    failed_pods = []
    for node_id, info in nodes.items():
        status = get_node_status(info["last_heartbeat"])
        print(f"Checking health for {node_id}: last_heartbeat={info['last_heartbeat']}, status={status}")
        if status == "failed":
            print(f"Node {node_id} detected as failed")
            failed_nodes.append(node_id)
            failed_pods.extend(info["pods"].keys())

    for node_id in failed_nodes:
        try:
            container = client.containers.get(node_id)
            container.stop()
            container.remove()
        except docker.errors.APIError:
            pass
        del nodes[node_id]

    for pod_id in failed_pods:
        suitable_node = next((node for node in nodes.values() if node["available_cores"] >= track_pods.get(pod_id, 2)), None)
        if suitable_node:
            schedule_pod(pod_id, suitable_node, track_pods.get(pod_id, 2))  # Use tracked cores or default to 2
        else:
            print(f"No suitable node found for pod {pod_id}")

    save_nodes()
    threading.Timer(10, check_node_health).start()

# Calculate cluster utilization for auto-scaling
def get_cluster_utilization():
    total_cores = sum(node["cpu_cores"] for node in nodes.values())
    available_cores = sum(node["available_cores"] for node in nodes.values())
    return (total_cores - available_cores) / total_cores if total_cores > 0 else 0

# Auto-scaling logic
def auto_scale():
    utilization = get_cluster_utilization()
    if utilization > (1 - SCALE_UP_THRESHOLD):
        print("Scaling up: Adding a new node")
        add_new_node(DEFAULT_NODE_CORES)
    elif utilization < SCALE_DOWN_THRESHOLD:
        print("Scaling down: Removing underutilized nodes")
        remove_underutilized_nodes()
    threading.Timer(60, auto_scale).start()  # Check every 60 seconds

# Add a new node for auto-scaling
def add_new_node(cpu_cores):
    node_id = generate_node_id()
    try:
        container = client.containers.run(
            "node-image",
            detach=True,
            environment={"NODE_ID": node_id, "API_SERVER_URL": "http://api-server:5000"},
            network="k8s-sim-net",
            name=node_id
        )
        nodes[node_id] = {
            "cpu_cores": cpu_cores,
            "available_cores": cpu_cores,
            "pods": {},
            "last_heartbeat": datetime.now(UTC).isoformat()
        }
        save_nodes()
        print(f"Auto-scaled node {node_id} added with {cpu_cores} CPU cores")
    except Exception as e:
        print(f"Failed to auto-scale node: {str(e)}")

# Remove underutilized nodes
def remove_underutilized_nodes():
    for node_id, node in list(nodes.items()):
        if node["available_cores"] > 0.5 * node["cpu_cores"] and not node["pods"]:
            try:
                container = client.containers.get(node_id)
                container.stop()
                container.remove()
                del nodes[node_id]
                save_nodes()
                print(f"Removed underutilized node {node_id}")
            except docker.errors.APIError:
                pass

# First-Fit scheduling algorithm with pod resource tracking
def schedule_pod(pod_id, target_node, required_cores):
    global track_pods
    if not target_node:
        print(f"No suitable node found for pod {pod_id}")
        return False
    if target_node["available_cores"] < required_cores:
        print(f"Insufficient cores on node for pod {pod_id}")
        return False
    target_node["pods"][pod_id] = required_cores
    target_node["available_cores"] -= required_cores
    track_pods[pod_id] = required_cores
    save_nodes()
    print(f"Scheduled pod {pod_id} on node {list(nodes.keys())[list(nodes.values()).index(target_node)]} with {required_cores} cores")
    return True

@app.route("/add_node", methods=["POST"])
def add_node():
    try:
        data = request.get_json()
        cpu_cores = data.get("cpu_cores", 0)
        if cpu_cores <= 0:
            return jsonify({"error": "Invalid CPU cores"}), 400

        node_id = generate_node_id()
        container = client.containers.run(
            "node-image",
            detach=True,
            environment={"NODE_ID": node_id, "API_SERVER_URL": "http://api-server:5000"},
            network="k8s-sim-net",
            name=node_id
        )

        nodes[node_id] = {
            "cpu_cores": cpu_cores,
            "available_cores": cpu_cores,
            "pods": {},
            "last_heartbeat": datetime.now(UTC).isoformat()
        }
        save_nodes()
        return jsonify({"message": f"Node {node_id} added with {cpu_cores} CPU cores"}), 201
    except Exception as e:
        return jsonify({"error": "Failed to add node", "details": str(e)}), 500

@app.route("/launch_pod", methods=["POST"])
def launch_pod():
    try:
        data = request.get_json()
        required_cores = data.get("cpu_cores", 2)
        if required_cores <= 0:
            return jsonify({"error": "Invalid CPU cores requirement"}), 400

        pod_id = generate_pod_id()
        suitable_node = next((node for node in nodes.values() if node["available_cores"] >= required_cores), None)
        if not suitable_node:
            return jsonify({"error": "Insufficient resources in cluster"}), 400

        if schedule_pod(pod_id, suitable_node, required_cores):
            return jsonify({
                "message": f"Pod {pod_id} launched successfully",
                "node_id": list(nodes.keys())[list(nodes.values()).index(suitable_node)],
                "pod_id": pod_id
            }), 201
        return jsonify({"error": "Failed to schedule pod"}), 500
    except Exception as e:
        return jsonify({"error": "Failed to launch pod", "details": str(e)}), 500

@app.route("/list_nodes", methods=["GET"])
def list_nodes():
    try:
        node_list = [
            {
                "node_id": node_id,
                "cpu_cores": info["cpu_cores"],
                "available_cores": info["available_cores"],
                "pods": list(info["pods"].keys()),
                "status": get_node_status(info["last_heartbeat"])
            }
            for node_id, info in nodes.items()
        ]
        return jsonify(node_list), 200
    except Exception as e:
        print(f"Error in list_nodes: {e}")
        return jsonify({"error": "Internal server error", "details": str(e)}), 500

@app.route("/heartbeat", methods=["POST"])
def heartbeat():
    print(f"Received request with Content-Type: {request.headers.get('Content-Type')}")
    try:
        data = request.get_json()
        if not data:
            return jsonify({"error": "No JSON data provided"}), 415
        node_id = data.get("node_id")
        print(f"Received heartbeat from {node_id}")
        if node_id in nodes:
            nodes[node_id]["last_heartbeat"] = datetime.now(UTC).isoformat()
            save_nodes()
            return jsonify({"message": f"Heartbeat received from {node_id}"}), 200
        return jsonify({"error": "Node not found"}), 404
    except Exception as e:
        return jsonify({"error": "Invalid JSON or server error", "details": str(e)}), 500

# Pod resource usage monitoring endpoint
@app.route("/pod_usage", methods=["GET"])
def pod_usage():
    try:
        pod_usage_list = []
        for node_id, node in nodes.items():
            for pod_id, cores in node["pods"].items():
                pod_usage_list.append({
                    "node_id": node_id,
                    "pod_id": pod_id,
                    "cpu_cores_allocated": cores
                })
        return jsonify(pod_usage_list), 200
    except Exception as e:
        return jsonify({"error": "Failed to retrieve pod usage", "details": str(e)}), 500

# Network policy management endpoints
@app.route("/add_policy", methods=["POST"])
def add_policy():
    try:
        data = request.get_json()
        source_pod = data.get("source_pod")
        target_pod = data.get("target_pod")
        action = data.get("action")
        if not source_pod or not target_pod or action not in ["allow", "deny"]:
            return jsonify({"error": "Invalid policy"}), 400
        policies.append({"source_pod": source_pod, "target_pod": target_pod, "action": action})
        return jsonify({"message": "Policy added"}), 201
    except Exception as e:
        return jsonify({"error": "Failed to add policy", "details": str(e)}), 500

@app.route("/list_policies", methods=["GET"])
def list_policies():
    return jsonify(policies), 200

@app.route("/communicate", methods=["POST"])
def communicate():
    try:
        data = request.get_json()
        source_pod = data.get("source_pod")
        target_pod = data.get("target_pod")
        for policy in policies:
            if policy["source_pod"] == source_pod and policy["target_pod"] == target_pod:
                if policy["action"] == "allow":
                    return jsonify({"message": "Communication allowed"}), 200
                else:
                    return jsonify({"message": "Communication denied"}), 403
        return jsonify({"message": "No policy found, default deny"}), 403
    except Exception as e:
        return jsonify({"error": "Failed to simulate communication", "details": str(e)}), 500

@app.route("/update_policy", methods=["POST"])
def update_policy():
    try:
        data = request.get_json()
        source_pod = data.get("source_pod")
        target_pod = data.get("target_pod")
        new_action = data.get("new_action")
        if not source_pod or not target_pod or new_action not in ["allow", "deny"]:
            return jsonify({"error": "Invalid policy update"}), 400
        for policy in policies:
            if policy["source_pod"] == source_pod and policy["target_pod"] == target_pod:
                policy["action"] = new_action
                return jsonify({"message": "Policy updated"}), 200
        return jsonify({"error": "Policy not found"}), 404
    except Exception as e:
        return jsonify({"error": "Failed to update policy", "details": str(e)}), 500


if __name__ == "__main__":
    load_nodes()
    threading.Timer(10, check_node_health).start()
    threading.Timer(60, auto_scale).start()  # Start auto-scaling
    app.run(host="0.0.0.0", port=5000)