优化代码复用

balance/base.go

@@ -7,7 +7,12 @@
 // Date : 2021-10-19 2:26 下午
 package balance
 
-import "sync"
+import (
+	"errors"
+	"sync"
+
+	"git.zhangdeman.cn/zhangdeman/gopkg/balance/define"
+)
 
 // base ...
 //
@@ -16,6 +21,7 @@ import "sync"
 // Date : 2:26 下午 2021/10/19
 type base struct {
 	lock      *sync.RWMutex
+	severList []*define.ServerNode
 }
 
 // Lock ...
@@ -53,3 +59,57 @@ func (b *base) RLock() {
 func (b *base) RUnlock() {
 	b.lock.RUnlock()
 }
+
+// AddServerNode 新添加一个服务器节点
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 10:07 下午 2021/10/20
+func (b *base) AddServerNode(hostIP string, port int) error {
+	if len(hostIP) == 0 || port <= 0 {
+		return errors.New("host ip or port is invalid")
+	}
+	b.Lock()
+	defer b.Unlock()
+	for _, item := range b.severList {
+		if item.HostIP == hostIP && item.Port == port {
+			return errors.New("host config is already exist")
+		}
+	}
+	b.severList = append(b.severList, &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 (b *base) Remove(hostIP string, port int, force bool) {
+	b.Lock()
+	defer b.Unlock()
+	tmpServerNode := make([]*define.ServerNode, 0)
+	for _, item := range b.severList {
+		if item.HostIP == hostIP && item.Port == port {
+			if force {
+				continue
+			}
+			item.Status = define.ServerNodeStatusRemove
+		}
+		tmpServerNode = append(tmpServerNode, item)
+	}
+	b.severList = tmpServerNode
+}
+
+// GetServerNodeList 获取服务器节点
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 10:33 下午 2021/10/20
+func (b *base) GetServerNodeList() []*define.ServerNode {
+	return b.severList
+}

balance/poll.go

@@ -20,7 +20,7 @@ import (
 // Author : go_developer@163.com<白茶清欢>
 //
 // Date : 12:40 下午 2021/10/19
-func NewPoll(serverList []*define.ServerNode) (*Poll, error) {
+func NewPoll(serverList []*define.ServerNode) (IBalance, error) {
 	if nil == serverList || len(serverList) == 0 {
 		return nil, errors.New("server list is empty")
 	}
@@ -28,7 +28,6 @@ func NewPoll(serverList []*define.ServerNode) (*Poll, error) {
 		base: base{
 			lock: &sync.RWMutex{},
 		},
-		serverList:         serverList,
 		currentServerIndex: 0,
 	}, nil
 }
@@ -40,7 +39,6 @@ func NewPoll(serverList []*define.ServerNode) (*Poll, error) {
 // Date : 12:41 下午 2021/10/19
 type Poll struct {
 	base
-	serverList         []*define.ServerNode
 	currentServerIndex int
 }
 
@@ -50,7 +48,8 @@ type Poll struct {
 //
 // Date : 12:43 下午 2021/10/19
 func (p *Poll) GetServerNode(req *http.Request) (*define.ServerNode, error) {
-	if len(p.serverList) == 0 {
+	severList := p.GetServerNodeList()
+	if len(severList) == 0 {
 		return nil, errors.New("server list is empty")
 	}
 	p.Lock()
@@ -60,10 +59,10 @@ func (p *Poll) GetServerNode(req *http.Request) (*define.ServerNode, error) {
 	)
 	// 循环次数
 	loopTimes := 0
-	for loopTimes < len(p.serverList) {
+	for loopTimes < len(severList) {
 		loopTimes++
-		p.currentServerIndex = (p.currentServerIndex + 1) % len(p.serverList)
+		p.currentServerIndex = (p.currentServerIndex + 1) % len(severList)
-		if serverNode = p.serverList[p.currentServerIndex]; serverNode.Status != define.ServerNodeStatusNormal {
+		if serverNode = severList[p.currentServerIndex]; serverNode.Status != define.ServerNodeStatusNormal {
 			continue
 		}
 		break
@@ -75,57 +74,3 @@ func (p *Poll) GetServerNode(req *http.Request) (*define.ServerNode, error) {
 
 	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
-}