feat: 添加过载保护、背压和熔断机制

internal/router/overload.go

@@ -0,0 +1,178 @@
+package router
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/edgeai/gateway/internal/config"
+)
+
+// OverloadStrategy defines how to handle requests when the primary model is overloaded.
+type OverloadStrategy int
+
+const (
+	StrategySameModelOtherInstance OverloadStrategy = iota // try same model on another instance
+	StrategySmallerLocalModel                              // fall back to a smaller local model
+	StrategyBackupEdgeNode                                 // route to a backup edge node
+	StrategyReject                                         // reject the request
+)
+
+// OverloadResolver resolves overload situations using configured strategies.
+type OverloadResolver struct {
+	strategies []OverloadStrategy
+	modelMap   *LogicalModelMapping
+}
+
+// NewOverloadResolver creates an overload resolver from routing config.
+func NewOverloadResolver(cfg *config.RoutingConfig, modelMap *LogicalModelMapping) *OverloadResolver {
+	r := &OverloadResolver{
+		modelMap: modelMap,
+	}
+	for _, s := range cfg.OverloadStrategy {
+		switch strings.ToLower(s) {
+		case "same_model_other_instance":
+			r.strategies = append(r.strategies, StrategySameModelOtherInstance)
+		case "smaller_local_model":
+			r.strategies = append(r.strategies, StrategySmallerLocalModel)
+		case "backup_edge_node":
+			r.strategies = append(r.strategies, StrategyBackupEdgeNode)
+		case "reject":
+			r.strategies = append(r.strategies, StrategyReject)
+		}
+	}
+	// Default: reject if no strategies configured
+	if len(r.strategies) == 0 {
+		r.strategies = []OverloadStrategy{StrategyReject}
+	}
+	return r
+}
+
+// OverloadResult contains the result of an overload resolution attempt.
+type OverloadResult struct {
+	Strategy  OverloadStrategy
+	Target    *ModelTarget // resolved fallback target, nil if rejected
+	Rejected  bool
+	Reason    string
+}
+
+// Resolve attempts to find a fallback target when the primary model is overloaded.
+// currentModel is the logical model that is overloaded.
+// excludeEndpoints is a set of endpoints already tried (to avoid loops).
+func (r *OverloadResolver) Resolve(currentModel string, excludeEndpoints map[string]bool) OverloadResult {
+	for _, strategy := range r.strategies {
+		switch strategy {
+		case StrategySameModelOtherInstance:
+			// Find same actual model on a different endpoint
+			target := r.findSameModelDifferentEndpoint(currentModel, excludeEndpoints)
+			if target != nil {
+				return OverloadResult{
+					Strategy: StrategySameModelOtherInstance,
+					Target:   target,
+				}
+			}
+		case StrategySmallerLocalModel:
+			// Find a model with smaller context window (proxy for "smaller")
+			target := r.findSmallerModel(currentModel, excludeEndpoints)
+			if target != nil {
+				return OverloadResult{
+					Strategy: StrategySmallerLocalModel,
+					Target:   target,
+				}
+			}
+		case StrategyBackupEdgeNode:
+			// Find any available model not yet tried
+			target := r.findAnyAvailable(excludeEndpoints)
+			if target != nil {
+				return OverloadResult{
+					Strategy: StrategyBackupEdgeNode,
+					Target:   target,
+				}
+			}
+		case StrategyReject:
+			return OverloadResult{
+				Strategy: StrategyReject,
+				Rejected:  true,
+				Reason:    "all overload strategies exhausted, request rejected",
+			}
+		}
+	}
+	return OverloadResult{
+		Strategy: StrategyReject,
+		Rejected:  true,
+		Reason:    fmt.Sprintf("no fallback available for model %s", currentModel),
+	}
+}
+
+func (r *OverloadResolver) findSameModelDifferentEndpoint(currentModel string, exclude map[string]bool) *ModelTarget {
+	r.modelMap.mu.RLock()
+	defer r.modelMap.mu.RUnlock()
+
+	current, ok := r.modelMap.mapping[currentModel]
+	if !ok {
+		return nil
+	}
+
+	for logical, target := range r.modelMap.mapping {
+		if logical == currentModel {
+			continue
+		}
+		if target.ActualModel == current.ActualModel && !exclude[target.Endpoint] {
+			return target
+		}
+	}
+	return nil
+}
+
+func (r *OverloadResolver) findSmallerModel(currentModel string, exclude map[string]bool) *ModelTarget {
+	r.modelMap.mu.RLock()
+	defer r.modelMap.mu.RUnlock()
+
+	current, ok := r.modelMap.mapping[currentModel]
+	if !ok {
+		return nil
+	}
+
+	var best *ModelTarget
+	for logical, target := range r.modelMap.mapping {
+		if logical == currentModel {
+			continue
+		}
+		if exclude[target.Endpoint] {
+			continue
+		}
+		// Pick model with smaller context window as "smaller" proxy
+		if target.ContextWindow < current.ContextWindow {
+			if best == nil || target.ContextWindow > best.ContextWindow {
+				best = target
+			}
+		}
+	}
+	return best
+}
+
+func (r *OverloadResolver) findAnyAvailable(exclude map[string]bool) *ModelTarget {
+	r.modelMap.mu.RLock()
+	defer r.modelMap.mu.RUnlock()
+
+	for _, target := range r.modelMap.mapping {
+		if !exclude[target.Endpoint] {
+			return target
+		}
+	}
+	return nil
+}
+
+// StrategyName returns a human-readable name for a strategy.
+func StrategyName(s OverloadStrategy) string {
+	switch s {
+	case StrategySameModelOtherInstance:
+		return "same_model_other_instance"
+	case StrategySmallerLocalModel:
+		return "smaller_local_model"
+	case StrategyBackupEdgeNode:
+		return "backup_edge_node"
+	case StrategyReject:
+		return "reject"
+	}
+	return "unknown"
+}