Merge pull request #16 from BadKid90s/loab_balance

feat(main): 增加LoadBalance负载均衡中间件

example/base/main.go

@@ -25,7 +25,10 @@ func main() {
 	//chain := gzip()
 
 	//cache middleware
-	chain := cache()
+	//chain := cache()
+
+	//cache middleware
+	chain := loadBalance()
 
 	//start serve
 	err := http.ListenAndServe(":8080", chain)
@@ -129,3 +132,33 @@ func cache() *knife.Chain {
 		})
 	return chain
 }
+
+func loadBalance() *knife.Chain {
+	nodes := []*middleware.ServiceNode{
+		{
+			Address: "127.0.0.1:8080",
+			Weight:  1,
+		},
+		{
+			Address: "127.0.0.2:8080",
+			Weight:  1,
+		},
+		{
+			Address: "127.0.0.3:8080",
+			Weight:  1,
+		},
+	}
+	chain := knife.NewChain().
+		Use(middleware.Logger()).
+		Use(middleware.Recover()).
+		Use(middleware.LoadBalanceProxy(middleware.LoadBalanceRandom, nodes)).
+		Use(func(context *knife.Context) {
+			data := "Hello World"
+			_, err := context.Writer.Write([]byte(data))
+			if err != nil {
+				panic(fmt.Sprintf("writer data error %s", err))
+			}
+			context.Abort(http.StatusOK)
+		})
+	return chain
+}

example/go.sum

@@ -0,0 +1,8 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

knife/go.mod

@@ -1,3 +1,10 @@
 module knife
 
 go 1.21
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/testify v1.8.4 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

knife/go.sum

@@ -0,0 +1,17 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

knife/middleware/load_balance_proxy.go

@@ -0,0 +1,139 @@
+package middleware
+
+import (
+	"fmt"
+	"hash/fnv"
+	"knife"
+	"math/rand"
+	"sync/atomic"
+	"time"
+)
+
+type LoadBalanceType int
+
+const (
+	LoadBalanceRandom           = iota //随机
+	LoadBalanceRoundRobin              //轮询
+	LoadBalanceWeightRoundRobin        //加权轮询
+	LoadBalanceHash                    //哈希
+)
+
+type ServiceNode struct {
+	Address string
+	Weight  int
+}
+
+type LoadBalancer struct {
+	nodes        []*ServiceNode
+	currentIndex int32
+}
+
+func LoadBalanceProxy(method LoadBalanceType, serviceNodes []*ServiceNode) knife.MiddlewareFunc {
+	lb := &LoadBalancer{
+		nodes:        append([]*ServiceNode{}, serviceNodes...),
+		currentIndex: 0,
+	}
+	switch method {
+	case LoadBalanceRandom:
+		return proxy(lb.random())
+	case LoadBalanceRoundRobin:
+		return proxy(lb.roundRobin())
+	case LoadBalanceWeightRoundRobin:
+		return proxy(lb.weightRoundRobin())
+	case LoadBalanceHash:
+		return proxyHash(lb)
+	default:
+		return proxy(lb.random())
+	}
+}
+
+// 随机算法的实现
+func (lb *LoadBalancer) random() string {
+	// 创建一个新的随机数种子
+	source := rand.NewSource(time.Now().UnixNano())
+	random := rand.New(source)
+	// 生成0到3之间的随机整数
+	randomNumber := random.Intn(len(lb.nodes))
+	//获取最终的服务地址
+	return lb.nodes[randomNumber].Address
+}
+
+// 轮询算法的实现
+func (lb *LoadBalancer) roundRobin() string {
+	index := atomic.AddInt32(&lb.currentIndex, 1)
+
+	i := index % int32(len(lb.nodes))
+
+	//获取最终的服务地址
+	return lb.nodes[i].Address
+
+}
+
+// 权重轮询算法的实现
+// 通过计算总权重并进行取模运算得到最终的服务节点下标
+func (lb *LoadBalancer) weightRoundRobin() string {
+	totalWeight := lb.getTotalWeight()
+
+	n := len(lb.nodes)
+	index := n - 1
+	hit := atomic.AddInt32(&lb.currentIndex, 1) % totalWeight
+
+	for i := 0; i < n; i++ {
+		weight := int32(lb.nodes[i].Weight)
+		hit = (hit + weight) % totalWeight
+		if hit < weight {
+			return lb.nodes[i].Address
+		}
+	}
+
+	//获取最终的服务地址
+	return lb.nodes[index].Address
+
+}
+
+// Hash算法实现
+// 通过请求地址计算hash值，让每次的请求都访问同一节点
+func (lb *LoadBalancer) hash(addr string) string {
+	// 创建一个 32 位的 FNV-1 哈希对象
+	hashed := fnv.New32()
+
+	// 对 int 类型的值 123 进行哈希计算
+	_, err := hashed.Write([]byte(addr))
+	if err != nil {
+		panic(fmt.Sprintf("compute hash value error：%s", err))
+	}
+	hashValue := hashed.Sum32()
+	// 输出哈希值
+	fmt.Println("哈希值为：", hashValue)
+
+	i := hashValue % uint32(len(lb.nodes))
+
+	//获取最终的服务地址
+	return lb.nodes[i].Address
+
+}
+
+// 获取所用的节点的权重
+func (lb *LoadBalancer) getTotalWeight() int32 {
+	totalWeight := 0
+	for _, node := range lb.nodes {
+		totalWeight += node.Weight
+	}
+	return int32(totalWeight)
+}
+
+func proxy(targetUrl string) knife.MiddlewareFunc {
+	return func(context *knife.Context) {
+		//代理请求
+		Proxy(targetUrl)
+	}
+}
+
+func proxyHash(balancer *LoadBalancer) knife.MiddlewareFunc {
+	return func(context *knife.Context) {
+		//通过ip计算得到最终的url
+		targetUrl := balancer.hash(context.Req.RemoteAddr)
+		//代理请求
+		Proxy(targetUrl)
+	}
+}

knife/middleware/load_balance_proxy_test.go

@@ -0,0 +1,114 @@
+package middleware
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var nodes = []*ServiceNode{
+	{
+		Address: "127.0.0.1:8080",
+		Weight:  1,
+	},
+	{
+		Address: "127.0.0.2:8080",
+		Weight:  1,
+	},
+	{
+		Address: "127.0.0.3:8080",
+		Weight:  1,
+	},
+}
+
+func TestLoadBalanceProxyHash(t *testing.T) {
+	lb := &LoadBalancer{
+		nodes:        nodes,
+		currentIndex: 0,
+	}
+
+	addr := "127.0.0.1:8080"
+	targetUrl1 := lb.hash(addr)
+	targetUrl2 := lb.hash(addr)
+	targetUrl3 := lb.hash(addr)
+	msg := "compute hash value no the same"
+	assert.Equal(t, targetUrl1, targetUrl2, msg)
+	assert.Equal(t, targetUrl1, targetUrl3, msg)
+	assert.Equal(t, targetUrl2, targetUrl3, msg)
+}
+
+func TestLoadBalanceProxyRandom(t *testing.T) {
+	lb := &LoadBalancer{
+		nodes:        nodes,
+		currentIndex: 0,
+	}
+
+	adders := []string{
+		"127.0.0.1:8080",
+		"127.0.0.2:8080",
+		"127.0.0.3:8080",
+	}
+
+	msg := "compute random index error"
+	for i := 0; i < 100; i++ {
+		targetUrl := lb.random()
+		assert.Contains(t, adders, targetUrl, msg)
+	}
+
+}
+
+func TestLoadBalanceProxyRoundRobin(t *testing.T) {
+	lb := &LoadBalancer{
+		nodes:        nodes,
+		currentIndex: 0,
+	}
+
+	adders := []string{
+		"127.0.0.1:8080",
+		"127.0.0.2:8080",
+		"127.0.0.3:8080",
+	}
+
+	msg := "compute roundRobin index error"
+	for i := 0; i < 3; i++ {
+		targetUrl := lb.roundRobin()
+		println(targetUrl)
+		index := (i + 1) % len(lb.nodes)
+		assert.Equal(t, adders[index], targetUrl, msg)
+	}
+}
+
+func TestLoadBalanceProxyWeightRoundRobin(t *testing.T) {
+	lb := &LoadBalancer{
+		nodes: []*ServiceNode{
+			{
+				Address: "127.0.0.1:8080",
+				Weight:  5,
+			},
+			{
+				Address: "127.0.0.2:8080",
+				Weight:  3,
+			},
+			{
+				Address: "127.0.0.3:8080",
+				Weight:  2,
+			},
+		},
+		currentIndex: 0,
+	}
+
+	addressNumbers := map[string]int{
+		"127.0.0.1:8080": 0,
+		"127.0.0.2:8080": 0,
+		"127.0.0.3:8080": 0,
+	}
+
+	msg := "compute weightRoundRobin index error"
+	for i := 0; i < 100; i++ {
+		targetUrl := lb.weightRoundRobin()
+		addressNumbers[targetUrl]++
+	}
+	assert.Equal(t, addressNumbers["127.0.0.1:8080"], 50, msg)
+	assert.Equal(t, addressNumbers["127.0.0.2:8080"], 30, msg)
+	assert.Equal(t, addressNumbers["127.0.0.3:8080"], 20, msg)
+
+}