load_balancer.py

from flask import Flask, redirect, request, jsonify
import requests
from threading import Lock
import logging
from datetime import datetime
import json
import os
import time

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('load_balancer.log'),
        logging.StreamHandler()
    ]
)

app = Flask(__name__)

# Configuration
SERVERS = [f"https://158.130.50.41:{port}" for port in range(7850, 7860)]  # Changed to HTTPS
HEALTH_CHECK_INTERVAL = 30  # seconds
MAX_FAILURES = 3

class LoadBalancer:
    def __init__(self):
        self.servers = SERVERS
        self.current_index = 0
        self.lock = Lock()
        self.server_status = {server: {'healthy': True, 'failures': 0} for server in self.servers}
        self.last_health_check = {server: datetime.now() for server in self.servers}

    def get_next_server(self):
        with self.lock:
            # Find the next healthy server
            start_index = self.current_index
            while True:
                server = self.servers[self.current_index]
                if self.server_status[server]['healthy']:
                    self.current_index = (self.current_index + 1) % len(self.servers)
                    return server
                
                self.current_index = (self.current_index + 1) % len(self.servers)
                if self.current_index == start_index:
                    # All servers are unhealthy
                    return None

    def mark_server_failure(self, server):
        with self.lock:
            self.server_status[server]['failures'] += 1
            if self.server_status[server]['failures'] >= MAX_FAILURES:
                self.server_status[server]['healthy'] = False
                logging.warning(f"Server {server} marked as unhealthy")

    def mark_server_success(self, server):
        with self.lock:
            self.server_status[server]['failures'] = 0
            self.server_status[server]['healthy'] = True

    def health_check(self, server):
        try:
            # Disable SSL verification for self-signed certificates
            response = requests.get(f"{server}/health", timeout=5, verify=False)
            if response.status_code == 200:
                self.mark_server_success(server)
                return True
            else:
                self.mark_server_failure(server)
                return False
        except Exception as e:
            logging.error(f"Health check failed for {server}: {str(e)}")
            self.mark_server_failure(server)
            return False

    def get_status(self):
        return {
            'servers': [
                {
                    'url': server,
                    'healthy': status['healthy'],
                    'failures': status['failures']
                }
                for server, status in self.server_status.items()
            ]
        }

# Initialize load balancer
load_balancer = LoadBalancer()

@app.route('/')
def proxy():
    server = load_balancer.get_next_server()
    if not server:
        return jsonify({'error': 'No healthy servers available'}), 503
    
    # Get the original request path and query parameters
    path = request.path
    query_string = request.query_string.decode('utf-8')
    
    # Construct the target URL
    target_url = f"{server}{path}"
    if query_string:
        target_url += f"?{query_string}"
    
    logging.info(f"Redirecting to {target_url}")
    return redirect(target_url)

@app.route('/health')
def health():
    return jsonify({'status': 'healthy'}), 200

@app.route('/status')
def status():
    return jsonify(load_balancer.get_status()), 200

def run_health_checks():
    """Background task to check server health"""
    while True:
        for server in SERVERS:
            load_balancer.health_check(server)
        time.sleep(HEALTH_CHECK_INTERVAL)

if __name__ == '__main__':
    # Start health check thread
    import threading
    health_check_thread = threading.Thread(target=run_health_checks, daemon=True)
    health_check_thread.start()
    
    # Run the Flask app on HTTP port 80
    app.run(host='0.0.0.0', port=7860)