main.py

from flask import Flask, request, Response, stream_with_context
import requests
import random
import os
import json
import time
import threading

app = Flask(__name__)

ROUTE_PREFIX = "/hf"

def get_backend_nodes():
    nodes_str = os.getenv("BACKEND_NODES", '[]')
    try:
        return json.loads(nodes_str)
    except json.JSONDecodeError:
        return []

NODES = get_backend_nodes()

# ==================== 全局禁用节点管理 ====================
DISABLED_NODES = {}
disabled_nodes_lock = threading.Lock()

def get_active_nodes():
    now = time.time()
    with disabled_nodes_lock:
        active_nodes = [
            node for node in NODES 
            if node not in DISABLED_NODES or DISABLED_NODES[node] <= now
        ]
        expired_nodes = [node for node, expire_time in DISABLED_NODES.items() if expire_time <= now]
        for node in expired_nodes:
            del DISABLED_NODES[node]
            
    return active_nodes
# ==========================================================

@app.route(f'{ROUTE_PREFIX}/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
def load_balancer(path):
    active_nodes = get_active_nodes()
    
    if not active_nodes:
        return "Error: No active backend nodes available at the moment.", 503

    available_nodes = list(active_nodes)
    req_json = request.get_json(silent=True) if request.method in ['POST', 'PUT', 'PATCH'] else None

    ESSENTIAL_HEADERS = {'content-type', 'accept', 'user-agent'}
    headers = {k: v for k, v in request.headers.items() if k.lower() in ESSENTIAL_HEADERS}

    while available_nodes:
        chosen_node = random.choice(available_nodes)
        target_url = f"{chosen_node.rstrip('/')}/{path}"
        
        try:
            print(f"[Info] Routing request to: {chosen_node}")
            response = requests.request(
                method=request.method,
                url=target_url,
                headers=headers,
                params=request.args,
                json=req_json,
                stream=True,
                timeout=120
            )
            
            # ==================== 核心状态码判定逻辑 ====================
            if response.status_code == 200:
                # 只有状态码为 200 时，才认为请求成功并进入流式传输
                def generate():
                    try:
                        for chunk in response.iter_content(chunk_size=8192):
                            yield chunk
                    except requests.exceptions.ChunkedEncodingError as e:
                        print(f"[Warning] Chunked stream ended prematurely from {chosen_node}: {e}")
                        return 
                    except Exception as e:
                        print(f"[Error] Unexpected error during streaming: {e}")
                        return

                proxy_response = Response(
                    stream_with_context(generate()),
                    status=response.status_code
                )

                for key, value in response.headers.items():
                    if key.lower() not in ('content-length', 'transfer-encoding', 'connection'):
                        proxy_response.headers[key] = value
                
                return proxy_response
            
            elif response.status_code == 429:
                # 【新加逻辑】429 状态码特判：不移除节点，保留在可用列表中，直接重试
                print(f"[Warning] Node {chosen_node} returned 429. Keeping it in pool and retrying...")
                time.sleep(0.5)  # 关键防护：稍微等待 0.5 秒，避免在极端情况下（如单节点）疯狂重试砸满 CPU
                continue
            
            else:
                # 状态码不是 200 且不是 429 的其他所有情况，均从当前请求的可用列表中移除
                if response.status_code == 403:
                    # 如果是 403 触发全局禁用机制
                    cooldown_seconds = 1 * 60
                    expire_time = time.time() + cooldown_seconds
                    with disabled_nodes_lock:
                        DISABLED_NODES[chosen_node] = expire_time
                    print(f"[Warning] Node {chosen_node} failed with 403. Disabling it globally for 1m.")
                else:
                    print(f"[Warning] Node {chosen_node} returned status {response.status_code}. Retrying another node...")
                
                available_nodes.remove(chosen_node)
                continue
            
        except requests.RequestException as e:
            # 捕获网络超时或连接断开等物理异常，依然移除并重试
            print(f"[Error] Network issue/Timeout with {chosen_node}: {e}. Retrying another node...")
            available_nodes.remove(chosen_node)

    return "Error: All available backend nodes failed (due to non-200 status or timeout).", 502


@app.route('/', methods=['GET'])
def index():
    return f"Load Balancer Running."

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)