github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/requester/downloader/loadbalance.go (about)

     1  package downloader
     2  
     3  import (
     4  	"net/http"
     5  	"sync/atomic"
     6  )
     7  
     8  type (
     9  	// LoadBalancerResponse 负载均衡响应状态
    10  	LoadBalancerResponse struct {
    11  		URL     string
    12  		Referer string
    13  	}
    14  
    15  	// LoadBalancerResponseList 负载均衡列表
    16  	LoadBalancerResponseList struct {
    17  		lbr    []*LoadBalancerResponse
    18  		cursor int32
    19  	}
    20  
    21  	LoadBalancerCompareFunc func(info map[string]string, subResp *http.Response) bool
    22  )
    23  
    24  // NewLoadBalancerResponseList 初始化负载均衡列表
    25  func NewLoadBalancerResponseList(lbr []*LoadBalancerResponse) *LoadBalancerResponseList {
    26  	return &LoadBalancerResponseList{
    27  		lbr: lbr,
    28  	}
    29  }
    30  
    31  // SequentialGet 顺序获取
    32  func (lbrl *LoadBalancerResponseList) SequentialGet() *LoadBalancerResponse {
    33  	if len(lbrl.lbr) == 0 {
    34  		return nil
    35  	}
    36  
    37  	if int(lbrl.cursor) >= len(lbrl.lbr) {
    38  		lbrl.cursor = 0
    39  	}
    40  
    41  	lbr := lbrl.lbr[int(lbrl.cursor)]
    42  	atomic.AddInt32(&lbrl.cursor, 1)
    43  	return lbr
    44  }
    45  
    46  // RandomGet 随机获取
    47  func (lbrl *LoadBalancerResponseList) RandomGet() *LoadBalancerResponse {
    48  	return lbrl.lbr[RandomNumber(0, len(lbrl.lbr))]
    49  }
    50  
    51  // AddLoadBalanceServer 增加负载均衡服务器
    52  func (der *Downloader) AddLoadBalanceServer(urls ...string) {
    53  	der.loadBalansers = append(der.loadBalansers, urls...)
    54  }
    55  
    56  // DefaultLoadBalancerCompareFunc 检测负载均衡的服务器是否一致
    57  func DefaultLoadBalancerCompareFunc(info map[string]string, subResp *http.Response) bool {
    58  	if info == nil || subResp == nil {
    59  		return false
    60  	}
    61  
    62  	for headerKey, value := range info {
    63  		if value != subResp.Header.Get(headerKey) {
    64  			return false
    65  		}
    66  	}
    67  
    68  	return true
    69  }