gopkg/balance/poll.go

// Package balance ...
//
// Description : 轮询机制的负载均衡
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-10-19 12:35 下午
package balance

import (
	"net/http"
	"sync"

	"git.zhangdeman.cn/zhangdeman/gopkg/balance/define"
	"github.com/pkg/errors"
)

// NewPoll 获取轮询实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 12:40 下午 2021/10/19
func NewPoll(serverList []*define.ServerNode) (*Poll, error) {
	if nil == serverList || len(serverList) == 0 {
		return nil, errors.New("server list is empty")
	}
	return &Poll{
		base: base{
			lock: &sync.RWMutex{},
		},
		serverList:         serverList,
		currentServerIndex: 0,
	}, nil
}

// Poll ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 12:41 下午 2021/10/19
type Poll struct {
	base
	serverList         []*define.ServerNode
	currentServerIndex int
}

// GetServerNode 获取服务节点
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 12:43 下午 2021/10/19
func (p *Poll) GetServerNode(req *http.Request) (*define.ServerNode, error) {
	if len(p.serverList) == 0 {
		return nil, errors.New("server list is empty")
	}
	p.Lock()
	defer p.Unlock()
	var (
		serverNode *define.ServerNode
	)
	// 循环次数
	loopTimes := 0
	for loopTimes < len(p.serverList) {
		loopTimes++
		p.currentServerIndex = (p.currentServerIndex + 1) % len(p.serverList)
		if serverNode = p.serverList[p.currentServerIndex]; serverNode.Status != define.ServerNodeStatusNormal {
			continue
		}
		break
	}

	if nil == serverNode || serverNode.Status != define.ServerNodeStatusNormal {
		return nil, errors.New("has no server node be used")
	}

	return serverNode, nil
}

// AddServerNode 新添加一个服务器节点
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:07 下午 2021/10/20
func (p *Poll) AddServerNode(hostIP string, port int) error {
	if len(hostIP) == 0 || port <= 0 {
		return errors.New("host ip or port is invalid")
	}
	p.Lock()
	defer p.Unlock()
	for _, item := range p.serverList {
		if item.HostIP == hostIP && item.Port == port {
			return errors.New("host config is already exist")
		}
	}
	p.serverList = append(p.serverList, &define.ServerNode{
		HostIP: hostIP,
		Port:   port,
		Status: define.ServerNodeStatusNormal,
	})
	return nil
}

// Remove 移除一个节点, force = true , 强制删除, force = false 逻辑删除, 设置状态
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:19 下午 2021/10/20
func (p *Poll) Remove(hostIP string, port int, force bool) {
	p.Lock()
	defer p.Unlock()
	tmpServerNode := make([]*define.ServerNode, 0)
	for _, item := range p.serverList {
		if item.HostIP == hostIP && item.Port == port {
			if force {
				continue
			}
			item.Status = define.ServerNodeStatusRemove
		}
		tmpServerNode = append(tmpServerNode, item)
	}
	p.serverList = tmpServerNode
}

// GetServerNodeList 获取服务器节点
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:33 下午 2021/10/20
func (p *Poll) GetServerNodeList() []*define.ServerNode {
	return p.serverList
}