Zoraxy/mod/dynamicproxy/loadbalance/loadbalance.go

package loadbalance

import (
	"sync"
	"time"

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

/*
	Load Balancer

	Handleing load balance request for upstream destinations
*/

type BalancePolicy int

// Data structure to hold the upstream server information
type Upstream struct {
	ProxyTargetIpOrDomain string   //Target IP address or domain name with port
	RequireTLS            bool     //Require TLS connection
	PreferedCountryCode   []string //Prefered country codes, default to empty string slice (not assigned)
	Priority              int      //Prirotiy of fallback, set all to 0 for round robin
}

// LoadBalanceRule stores the load balance option of this host. If there are no usable upstreams for lb, the primary one will be used
type LoadBalanceRule struct {
	Upstreams        []Upstream //Reverse proxy upstream servers
	UseStickySession bool       //Enable stick session
}

// PrimaryRoutingRule stores the primary routing target of this host
type PrimaryRoutingRule struct {
	MatchingDomainOrIp string //Primary proxy upstream origin server
	RequireTLS         bool   //Target require TLS
}

type Options struct {
	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
}

// 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,
	}
}

func (m *RouteManager) UpdateKeepAliveTargets(pingTargets []*PrimaryRoutingRule) {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	go func() {
		for {
			select {
			case <-m.onlineStatusTickerStop:
				ticker.Stop()
				return
			case <-ticker.C:
				for _, target := range pingTargets {
					isOnline := PingTarget(target)
					m.LoadBalanceMap.Store(target.MatchingDomainOrIp, isOnline)
				}
			}
		}
	}()
}

// GetProxyTargetIP Get the proxy target given the primary routing rule and load balance options
func (m *RouteManager) GetProxyTargetIP(pr *PrimaryRoutingRule, lbr *LoadBalanceRule) {

}

func (m *RouteManager) Close() {
	m.onlineStatusTickerStop <- true
}

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