adding load balancing to proxy

fast-server/config/config.go

@@ -14,11 +14,18 @@ import (
 
 var ProductionConfigPath = "/etc/fast/config.yaml"
 
+const (
+	LoadBalanceMethodRoundRobin = "round_robin"
+	LoadBalanceMethodLeastConn  = "least_conn"
+)
+
 type ProxyConfig struct {
-	Host               string `yaml:"host"`
-	Port               int    `yaml:"port"`
-	Protocol           string `yaml:"protocol,omitempty"`
-	InsecureSkipVerify bool   `yaml:"insecure_skip_verify,omitempty"`
+	Host               string   `yaml:"host,omitempty"`
+	Hosts              []string `yaml:"hosts,omitempty"`
+	Port               int      `yaml:"port"`
+	Protocol           string   `yaml:"protocol,omitempty"`
+	InsecureSkipVerify bool     `yaml:"insecure_skip_verify,omitempty"`
+	LoadBalanceMethod  string   `yaml:"load_balance_method,omitempty"`
 }
 
 func (p *ProxyConfig) setDefaults() {
@@ -29,15 +36,23 @@ func (p *ProxyConfig) setDefaults() {
 }
 
 func (p *ProxyConfig) validate() error {
-	if p.Host == "" {
-		return fmt.Errorf("proxy host cannot be empty")
+	if p.Host == "" && len(p.Hosts) == 0 {
+		return fmt.Errorf("proxy host or hosts must be specified")
 	}
 	if p.Port <= 0 || p.Port > 65535 {
 		return fmt.Errorf("invalid proxy port: %d", p.Port)
 	}
 	if p.Protocol != "http" && p.Protocol != "https" {
 		return fmt.Errorf("invalid proxy protocol: %s", p.Protocol)
 	}
+	if p.LoadBalanceMethod != "" &&
+		p.LoadBalanceMethod != LoadBalanceMethodRoundRobin &&
+		p.LoadBalanceMethod != LoadBalanceMethodLeastConn {
+		return fmt.Errorf("invalid load balance method: %s, must be one of: %s, %s",
+			p.LoadBalanceMethod,
+			LoadBalanceMethodRoundRobin,
+			LoadBalanceMethodLeastConn)
+	}
 	return nil
 }
 

fast-server/handlers/proxy_handler.go

@@ -6,13 +6,65 @@ import (
 	"fmt"
 	"github.com/labstack/echo/v4"
 	"github.com/labstack/echo/v4/middleware"
+	"math"
 	"net/http"
 	"net/url"
 	"sort"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 )
 
+var (
+	loadBalancerCounter uint32
+	activeConnections   sync.Map
+)
+
+const (
+	LoadBalanceMethodRoundRobin = "round_robin"
+	LoadBalanceMethodLeastConn  = "least_conn"
+)
+
+func getNextHost(proxyConfig config.ProxyConfig) string {
+	// If only single host is specified, return it
+	if len(proxyConfig.Hosts) == 0 {
+		return proxyConfig.Host
+	}
+
+	// Choose load balancing method
+	switch proxyConfig.LoadBalanceMethod {
+	case LoadBalanceMethodLeastConn:
+		return getLeastConnHost(proxyConfig.Hosts)
+	case LoadBalanceMethodRoundRobin:
+		return getRoundRobinHost(proxyConfig.Hosts)
+	default:
+		// If no method specified or unknown, default to round_robin
+		return getRoundRobinHost(proxyConfig.Hosts)
+	}
+}
+
+func getRoundRobinHost(hosts []string) string {
+	next := atomic.AddUint32(&loadBalancerCounter, 1)
+	return hosts[next%uint32(len(hosts))]
+}
+
+func getLeastConnHost(hosts []string) string {
+	var selectedHost string
+	minConns := int64(math.MaxInt64)
+
+	for _, host := range hosts {
+		conns, _ := activeConnections.LoadOrStore(host, int64(0))
+		currentConns := conns.(int64)
+		if currentConns < minConns {
+			minConns = currentConns
+			selectedHost = host
+		}
+	}
+
+	return selectedHost
+}
+
 func getMatchingLocation(path string, locations []config.Location) *config.Location {
 	// Sort locations by path length in descending order (longest first)
 	sortedLocations := make([]config.Location, len(locations))
@@ -42,16 +94,30 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
 		return echo.ErrNotFound
 	}
 
-	c.Logger().Infof("Matched location path: %s", location.Path)
+	c.Logger().Debugf("Matched location path: %s", location.Path)
 
 	// Use location's proxy config
 	proxyConfig := location.Proxy
+
+	// Get next host using load balancing
+	selectedHost := getNextHost(proxyConfig)
+
+	// Increment active connections for the selected host if using least_conn
+	if proxyConfig.LoadBalanceMethod == LoadBalanceMethodLeastConn {
+		var conns int64
+		actual, _ := activeConnections.LoadOrStore(selectedHost, &conns)
+		connPtr := actual.(*int64)
+		atomic.AddInt64(connPtr, 1)
+		defer atomic.AddInt64(connPtr, -1)
+		c.Logger().Debugf("Active connections for host %s: %d", selectedHost, atomic.LoadInt64(connPtr))
+	}
+
 	targetScheme := proxyConfig.Protocol
 	if targetScheme == "" {
 		targetScheme = "http"
 	}
 
-	target, err := url.Parse(fmt.Sprintf("%s://%s:%d", targetScheme, proxyConfig.Host, proxyConfig.Port))
+	target, err := url.Parse(fmt.Sprintf("%s://%s:%d", targetScheme, selectedHost, proxyConfig.Port))
 	if err != nil {
 		c.Logger().Errorf("Error parsing proxy URL: %v", err)
 		return echo.ErrInternalServerError
@@ -71,9 +137,17 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
 		}
 	}
 
-	c.Logger().Infof("Using rewrite map: %v", rewriteMap)
-	c.Logger().Infof("Proxying %s request from %s to %s%s",
-		c.Request().Method, originalHost, target.String(), requestPath)
+	c.Logger().Debugf("Using rewrite map: %v", rewriteMap)
+
+	// More detailed logging for load balancing
+	c.Logger().Infof("[LoadBalancer] %s request: %s → %s (host: %s, method: %s, scheme: %s)",
+		proxyConfig.LoadBalanceMethod,
+		originalHost+requestPath,
+		target.String(),
+		selectedHost,
+		c.Request().Method,
+		targetScheme,
+	)
 
 	transport := &http.Transport{
 		Proxy: http.ProxyFromEnvironment,
@@ -113,17 +187,23 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
 		Rewrite:   rewriteMap,
 		Transport: transport,
 		ModifyResponse: func(res *http.Response) error {
-			c.Logger().Infof("Proxy response status: %d for path: %s", res.StatusCode, requestPath)
+			c.Logger().Infof("[LoadBalancer] Response from %s: %d %s",
+				selectedHost,
+				res.StatusCode,
+				http.StatusText(res.StatusCode),
+			)
 			return nil
 		},
 	})
 
-	// Set proxy headers
+	// Set proxy headers with more detailed tracking
 	c.Request().Header.Set("X-Forwarded-Host", originalHost)
 	c.Request().Header.Set("X-Real-IP", c.RealIP())
 	c.Request().Header.Set("X-Forwarded-For", c.RealIP())
 	c.Request().Header.Set("X-Forwarded-Proto", sourceScheme)
 	c.Request().Header.Set("X-Original-URI", requestPath)
+	c.Request().Header.Set("X-Load-Balanced-Host", selectedHost)
+	c.Request().Header.Set("X-Load-Balance-Method", proxyConfig.LoadBalanceMethod)
 	c.Request().Host = originalHost
 
 	return proxyMiddleware(func(c echo.Context) error {