From 964afa3c58f7f5c79564e3bf4bb19a54ef8fe815 Mon Sep 17 00:00:00 2001
From: maxing <maxing@njust.edu.cn>
Date: Sun, 9 Oct 2022 16:19:45 +0800
Subject: [PATCH] refactor: update traffic filter to support both header and
 weight config.

---
 pixiu/pkg/filter/traffic/models.go  |  4 +--
 pixiu/pkg/filter/traffic/traffic.go | 45 +++++++++++++++++------------
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/pixiu/pkg/filter/traffic/models.go b/pixiu/pkg/filter/traffic/models.go
index dbdd37cf0..eb6025df7 100644
--- a/pixiu/pkg/filter/traffic/models.go
+++ b/pixiu/pkg/filter/traffic/models.go
@@ -21,10 +21,10 @@ import (
 	"net/http"
 )
 
-func headerSpilt(req *http.Request, value string) bool {
+func spiltHeader(req *http.Request, value string) bool {
 	return req.Header.Get(string(canaryByHeader)) == value
 }
 
-func weightSpilt(weight, floor, ceil int) bool {
+func spiltWeight(weight, floor, ceil int) bool {
 	return weight >= floor && weight < ceil
 }
diff --git a/pixiu/pkg/filter/traffic/traffic.go b/pixiu/pkg/filter/traffic/traffic.go
index 46f95e188..3fd82823c 100644
--- a/pixiu/pkg/filter/traffic/traffic.go
+++ b/pixiu/pkg/filter/traffic/traffic.go
@@ -18,9 +18,7 @@
 package traffic
 
 import (
-	"fmt"
 	"math/rand"
-	"strconv"
 	"strings"
 	"time"
 )
@@ -72,7 +70,7 @@ type (
 		Name           string `yaml:"name" json:"name" mapstructure:"name"`
 		Router         string `yaml:"router" json:"router" mapstructure:"router"`
 		CanaryByHeader string `yaml:"canary-by-header" json:"canary-by-header" mapstructure:"canary-by-header"`
-		CanaryWeight   string `yaml:"canary-weight" json:"canary-weight" mapstructure:"canary-weight"`
+		CanaryWeight   int    `yaml:"canary-weight" json:"canary-weight" mapstructure:"canary-weight"`
 	}
 )
 
@@ -105,33 +103,45 @@ func (factory *FilterFactory) PrepareFilterChain(ctx *http.HttpContext, chain fi
 }
 
 func (f *Filter) Decode(ctx *http.HttpContext) filter.FilterStatus {
+	cluster := ""
 	if f.Rules != nil {
 		for _, wp := range f.Rules {
-			if f.traffic(wp, ctx) {
-				ctx.Route.Cluster = wp.Cluster.Name
+			if f.trafficHeader(wp, ctx) {
+				cluster = wp.Cluster.Name
 				logger.Debugf("[dubbo-go-pixiu] execute traffic split to cluster %s", wp.Cluster.Name)
 				break
 			}
 		}
+		if cluster == "" {
+			for _, wp := range f.Rules {
+				if f.trafficWeight(wp, ctx) {
+					ctx.Route.Cluster = wp.Cluster.Name
+					cluster = wp.Cluster.Name
+					logger.Debugf("[dubbo-go-pixiu] execute traffic split to cluster %s", wp.Cluster.Name)
+					break
+				}
+			}
+		}
+		if cluster != "" {
+			ctx.Route.Cluster = cluster
+		}
 	} else {
 		logger.Warnf("[dubbo-go-pixiu] execute traffic split fail because of empty rules.")
 	}
 	return filter.Continue
 }
 
-func (f *Filter) traffic(c *ClusterWrapper, ctx *http.HttpContext) bool {
+func (f *Filter) trafficHeader(c *ClusterWrapper, ctx *http.HttpContext) bool {
+	return spiltHeader(ctx.Request, c.header)
+}
+
+func (f *Filter) trafficWeight(c *ClusterWrapper, ctx *http.HttpContext) bool {
 	if f.weight == unInitialize {
 		rand.Seed(time.Now().UnixNano())
 		f.weight = rand.Intn(100) + 1
 	}
 
-	res := false
-	if c.header != "" {
-		res = headerSpilt(ctx.Request, c.header)
-	} else if !res && c.weightFloor != -1 && c.weightCeil != -1 {
-		res = weightSpilt(f.weight, c.weightFloor, c.weightCeil)
-	}
-	return res
+	return spiltWeight(f.weight, c.weightFloor, c.weightCeil)
 }
 
 func (factory *FilterFactory) rulesMatch(f *Filter, path string) []*ClusterWrapper {
@@ -156,15 +166,12 @@ func (factory *FilterFactory) rulesMatch(f *Filter, path string) []*ClusterWrapp
 						wp.header = cluster.CanaryByHeader
 					}
 				}
-				if cluster.CanaryWeight != "" {
-					val, err := strconv.Atoi(cluster.CanaryWeight)
-					if err != nil || val <= 0 {
-						logger.Errorf(fmt.Sprintf("Wrong canary-weight value: %v", cluster.CanaryWeight))
-					}
+				if cluster.CanaryWeight > 0 && cluster.CanaryWeight <= 100 {
 					wp.weightFloor = up
-					up += val
+					up += cluster.CanaryWeight
 					if up > 100 {
 						logger.Errorf("[dubbo-go-pixiu] clusters' weight sum more than 100 in %v service!", cluster.Router)
+						continue
 					}
 					wp.weightCeil = up
 				}