Zoraxy/mod/dynamicproxy/loadbalance/loadbalance.go

package loadbalance

import (
	"strings"
	"sync"
	"sync/atomic"

	"imuslab.com/zoraxy/mod/dynamicproxy/dpcore"
	"imuslab.com/zoraxy/mod/geodb"
	"imuslab.com/zoraxy/mod/info/logger"
)

/*
	Load Balancer

	Handleing load balance request for upstream destinations
*/

type Options struct {
	UseActiveHealthCheck bool         //Use active health check, default to false
	Geodb                *geodb.Store //GeoIP resolver for checking incoming request origin country
	Logger               *logger.Logger
}

type RouteManager struct {
	LoadBalanceMap         sync.Map  //Sync map to store the last load balance state of a given node
	OnlineStatusMap        sync.Map  //Sync map to store the online status of a given ip address or domain name
	onlineStatusTickerStop chan bool //Stopping channel for the online status pinger
	Options                Options   //Options for the load balancer
}

/* Upstream or Origin Server */
type Upstream struct {
	//Upstream Proxy Configs
	OriginIpOrDomain         string //Target IP address or domain name with port
	RequireTLS               bool   //Require TLS connection
	SkipCertValidations      bool   //Set to true to accept self signed certs
	SkipWebSocketOriginCheck bool   //Skip origin check on websocket upgrade connections

	//Load balancing configs
	Priority int  //Prirotiy of fallback, set all to 0 for round robin
	MaxConn  int  //Maxmium connection to this server, 0 for unlimited
	Disabled bool //If this upstream is disabled

	currentConnectionCounts atomic.Uint64 //Counter for number of client currently connected
	proxy                   *dpcore.ReverseProxy
}

// Create a new load balancer
func NewLoadBalancer(options *Options) *RouteManager {
	onlineStatusCheckerStopChan := make(chan bool)

	return &RouteManager{
		LoadBalanceMap:         sync.Map{},
		OnlineStatusMap:        sync.Map{},
		onlineStatusTickerStop: onlineStatusCheckerStopChan,
		Options:                *options,
	}
}

// UpstreamsReady checks if the group of upstreams contains at least one
// origin server that is ready
func (m *RouteManager) UpstreamsReady(upstreams []*Upstream) bool {
	for _, upstream := range upstreams {
		if upstream.Disabled {
			//This upstream is disabled. Assume offline
			continue
		}
		if upstream.IsReady() {
			return true
		}
	}
	return false
}

// String format and convert a list of upstream into a string representations
func GetUpstreamsAsString(upstreams []*Upstream) string {
	targets := []string{}
	for _, upstream := range upstreams {
		targets = append(targets, upstream.String())
	}
	return strings.Join(targets, ", ")
}

func (m *RouteManager) Close() {
	if m.onlineStatusTickerStop != nil {
		m.onlineStatusTickerStop <- true
	}

}

// Print debug message
func (m *RouteManager) debugPrint(message string, err error) {
	m.Options.Logger.PrintAndLog("LoadBalancer", message, err)
}