Merge pull request #112 from realab/bm/criticality

criticality and client criticality with bm

pkg/net/criticality/criticality.go

@@ -0,0 +1,57 @@
+package criticality
+
+// Criticality is
+type Criticality string
+
+// criticality
+var (
+	// EmptyCriticality is used to mark any invalid criticality, and the empty criticality will be parsed as the default criticality later.
+	EmptyCriticality = Criticality("")
+	// CriticalPlus is reserved for the most critical requests, those that will result in serious user-visible impact if they fail.
+	CriticalPlus = Criticality("CRITICAL_PLUS")
+	// Critical is the default value for requests sent from production jobs. These requests will result in user-visible impact, but the impact may be less severe than those of CRITICAL_PLUS. Services are expected to provision enough capacity for all expected CRITICAL and CRITICAL_PLUS traffic.
+	Critical = Criticality("CRITICAL")
+	// SheddablePlus is traffic for which partial unavailability is expected. This is the default for batch jobs, which can retry requests minutes or even hours later.
+	SheddablePlus = Criticality("SHEDDABLE_PLUS")
+	// Sheddable is traffic for which frequent partial unavailability and occasional full unavailability is expected.
+	Sheddable = Criticality("SHEDDABLE")
+
+	// higher is more critical
+	_criticalityEnum = map[Criticality]int{
+		CriticalPlus:  40,
+		Critical:      30,
+		SheddablePlus: 20,
+		Sheddable:     10,
+	}
+
+	_defaultCriticality = Critical
+)
+
+// Value is used to get criticality value, higher value is more critical.
+func Value(in Criticality) int {
+	v, ok := _criticalityEnum[in]
+	if !ok {
+		return _criticalityEnum[_defaultCriticality]
+	}
+	return v
+}
+
+// Higher will compare the input criticality with self, return true if the input is more critical than self.
+func (c Criticality) Higher(in Criticality) bool {
+	return Value(in) > Value(c)
+}
+
+// Parse will parse raw criticality string as valid critality. Any invalid input will parse as empty criticality.
+func Parse(raw string) Criticality {
+	crtl := Criticality(raw)
+	if _, ok := _criticalityEnum[crtl]; ok {
+		return crtl
+	}
+	return EmptyCriticality
+}
+
+// Exist is used to check criticality is exist in several enumeration.
+func Exist(c Criticality) bool {
+	_, ok := _criticalityEnum[c]
+	return ok
+}

pkg/net/http/blademaster/client.go

@@ -253,9 +253,11 @@ func (client *Client) Raw(c context.Context, req *xhttp.Request, v ...string) (b
 	setTimeout(req, timeout)
 	req = req.WithContext(c)
 	setCaller(req)
-	if color := metadata.String(c, metadata.Color); color != "" {
-		setColor(req, color)
-	}
+	metadata.Range(c,
+		func(key string, value interface{}) {
+			setMetadata(req, key, value)
+		},
+		metadata.IsOutgoingKey)
 	if resp, err = client.client.Do(req); err != nil {
 		err = pkgerr.Wrapf(err, "host:%s, url:%s", req.URL.Host, realURL(req))
 		code = "failed"

pkg/net/http/blademaster/criticality.go

@@ -0,0 +1,21 @@
+package blademaster
+
+import (
+	criticalityPkg "github.com/bilibili/kratos/pkg/net/criticality"
+	"github.com/bilibili/kratos/pkg/net/metadata"
+
+	"github.com/pkg/errors"
+)
+
+// Criticality is
+func Criticality(pathCriticality criticalityPkg.Criticality) HandlerFunc {
+	if !criticalityPkg.Exist(pathCriticality) {
+		panic(errors.Errorf("This criticality is not exist: %s", pathCriticality))
+	}
+	return func(ctx *Context) {
+		md, ok := metadata.FromContext(ctx)
+		if ok {
+			md[metadata.Criticality] = string(pathCriticality)
+		}
+	}
+}

pkg/net/http/blademaster/metadata.go

@@ -1,6 +1,7 @@
 package blademaster
 
 import (
+	"fmt"
 	"net/http"
 	"strconv"
 	"strings"
@@ -8,6 +9,8 @@ import (
 
 	"github.com/bilibili/kratos/pkg/conf/env"
 	"github.com/bilibili/kratos/pkg/log"
+	"github.com/bilibili/kratos/pkg/net/criticality"
+	"github.com/bilibili/kratos/pkg/net/metadata"
 
 	"github.com/pkg/errors"
 )
@@ -15,52 +18,64 @@ import (
 const (
 	// http head
 	_httpHeaderUser         = "x1-bmspy-user"
-	_httpHeaderColor        = "x1-bmspy-color"
 	_httpHeaderTimeout      = "x1-bmspy-timeout"
-	_httpHeaderMirror       = "x1-bmspy-mirror"
 	_httpHeaderRemoteIP     = "x-backend-bm-real-ip"
 	_httpHeaderRemoteIPPort = "x-backend-bm-real-ipport"
 )
 
-// mirror return true if x-bmspy-mirror in http header and its value is 1 or true.
-func mirror(req *http.Request) bool {
-	mirrorStr := req.Header.Get(_httpHeaderMirror)
-	if mirrorStr == "" {
-		return false
-	}
-	val, err := strconv.ParseBool(mirrorStr)
-	if err != nil {
-		log.Warn("blademaster: failed to parse mirror: %+v", errors.Wrap(err, mirrorStr))
-		return false
-	}
-	if !val {
-		log.Warn("blademaster: request mirrorStr value :%s is false", mirrorStr)
-	}
-	return val
-}
+const (
+	_httpHeaderMetadata = "x-bm-metadata-"
+)
 
-// setCaller set caller into http request.
-func setCaller(req *http.Request) {
-	req.Header.Set(_httpHeaderUser, env.AppID)
+var _parser = map[string]func(string) interface{}{
+	"mirror": func(mirrorStr string) interface{} {
+		if mirrorStr == "" {
+			return false
+		}
+		val, err := strconv.ParseBool(mirrorStr)
+		if err != nil {
+			log.Warn("blademaster: failed to parse mirror: %+v", errors.Wrap(err, mirrorStr))
+			return false
+		}
+		if !val {
+			log.Warn("blademaster: request mirrorStr value :%s is false", mirrorStr)
+		}
+		return val
+	},
+	"criticality": func(in string) interface{} {
+		if crtl := criticality.Criticality(in); crtl != criticality.EmptyCriticality {
+			return string(crtl)
+		}
+		return string(criticality.Critical)
+	},
 }
 
-// caller get caller from http request.
-func caller(req *http.Request) string {
-	return req.Header.Get(_httpHeaderUser)
+func parseMetadataTo(req *http.Request, to metadata.MD) {
+	for rawKey := range req.Header {
+		key := strings.ReplaceAll(strings.TrimLeft(strings.ToLower(rawKey), _httpHeaderMetadata), "-", "_")
+		rawValue := req.Header.Get(rawKey)
+		var value interface{} = rawValue
+		parser, ok := _parser[key]
+		if ok {
+			value = parser(rawValue)
+		}
+		to[key] = value
+	}
+	return
 }
 
-// setColor set color into http request.
-func setColor(req *http.Request, color string) {
-	req.Header.Set(_httpHeaderColor, color)
+func setMetadata(req *http.Request, key string, value interface{}) {
+	strV, ok := value.(string)
+	if !ok {
+		return
+	}
+	header := fmt.Sprintf("%s%s", _httpHeaderMetadata, strings.ReplaceAll(key, "_", "-"))
+	req.Header.Set(header, strV)
 }
 
-// color get color from http request.
-func color(req *http.Request) string {
-	c := req.Header.Get(_httpHeaderColor)
-	if c == "" {
-		c = env.Color
-	}
-	return c
+// setCaller set caller into http request.
+func setCaller(req *http.Request) {
+	req.Header.Set(_httpHeaderUser, env.AppID)
 }
 
 // setTimeout set timeout into http request.

pkg/net/http/blademaster/server.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/bilibili/kratos/pkg/conf/dsn"
 	"github.com/bilibili/kratos/pkg/log"
+	"github.com/bilibili/kratos/pkg/net/criticality"
 	"github.com/bilibili/kratos/pkg/net/ip"
 	"github.com/bilibili/kratos/pkg/net/metadata"
 	"github.com/bilibili/kratos/pkg/stat"
@@ -261,12 +262,11 @@ func (engine *Engine) handleContext(c *Context) {
 		tm = ctm
 	}
 	md := metadata.MD{
-		metadata.Color:      color(req),
-		metadata.RemoteIP:   remoteIP(req),
-		metadata.RemotePort: remotePort(req),
-		metadata.Caller:     caller(req),
-		metadata.Mirror:     mirror(req),
+		metadata.RemoteIP:    remoteIP(req),
+		metadata.RemotePort:  remotePort(req),
+		metadata.Criticality: string(criticality.Critical),
 	}
+	parseMetadataTo(req, md)
 	ctx := metadata.NewContext(context.Background(), md)
 	if tm > 0 {
 		c.Context, cancel = context.WithTimeout(ctx, tm)

pkg/net/http/blademaster/server_test.go

@@ -0,0 +1,137 @@
+package blademaster
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	criticalityPkg "github.com/bilibili/kratos/pkg/net/criticality"
+	"github.com/bilibili/kratos/pkg/net/metadata"
+	xtime "github.com/bilibili/kratos/pkg/time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	sonce sync.Once
+
+	SockAddr  = "localhost:18080"
+	curEngine atomic.Value
+)
+
+func uri(base, path string) string {
+	return fmt.Sprintf("%s://%s%s", "http", base, path)
+}
+
+func shutdown() {
+	if err := curEngine.Load().(*Engine).Shutdown(context.TODO()); err != nil {
+		panic(err)
+	}
+}
+
+func setupHandler(engine *Engine) {
+	// set the global timeout is 2 second
+	engine.conf.Timeout = xtime.Duration(time.Second * 2)
+
+	engine.Ping(func(ctx *Context) {
+		ctx.AbortWithStatus(200)
+	})
+
+	engine.GET("/criticality/api", Criticality(criticalityPkg.Critical), func(ctx *Context) {
+		ctx.String(200, "%s", metadata.String(ctx, metadata.Criticality))
+	})
+	engine.GET("/criticality/none/api", func(ctx *Context) {
+		ctx.String(200, "%s", metadata.String(ctx, metadata.Criticality))
+	})
+}
+
+func startServer() {
+	e := DefaultServer(nil)
+	setupHandler(e)
+	go e.Run(SockAddr)
+	curEngine.Store(e)
+	time.Sleep(time.Second)
+}
+
+func TestCriticality(t *testing.T) {
+	startServer()
+	defer shutdown()
+
+	tests := []*struct {
+		path     string
+		crtl     criticalityPkg.Criticality
+		expected criticalityPkg.Criticality
+	}{
+		{
+			"/criticality/api",
+			criticalityPkg.EmptyCriticality,
+			criticalityPkg.Critical,
+		},
+		{
+			"/criticality/api",
+			criticalityPkg.CriticalPlus,
+			criticalityPkg.Critical,
+		},
+		{
+			"/criticality/api",
+			criticalityPkg.SheddablePlus,
+			criticalityPkg.Critical,
+		},
+	}
+	client := &http.Client{}
+	for _, testCase := range tests {
+		req, err := http.NewRequest("GET", uri(SockAddr, testCase.path), nil)
+		assert.NoError(t, err)
+		req.Header.Set("x-bm-metadata-criticality", string(testCase.crtl))
+		resp, err := client.Do(req)
+		assert.NoError(t, err)
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, testCase.expected, criticalityPkg.Criticality(body))
+	}
+}
+
+func TestNoneCriticality(t *testing.T) {
+	startServer()
+	defer shutdown()
+
+	tests := []*struct {
+		path     string
+		crtl     criticalityPkg.Criticality
+		expected criticalityPkg.Criticality
+	}{
+		{
+			"/criticality/none/api",
+			criticalityPkg.EmptyCriticality,
+			criticalityPkg.Critical,
+		},
+		{
+			"/criticality/none/api",
+			criticalityPkg.CriticalPlus,
+			criticalityPkg.CriticalPlus,
+		},
+		{
+			"/criticality/none/api",
+			criticalityPkg.SheddablePlus,
+			criticalityPkg.SheddablePlus,
+		},
+	}
+	client := &http.Client{}
+	for _, testCase := range tests {
+		req, err := http.NewRequest("GET", uri(SockAddr, testCase.path), nil)
+		assert.NoError(t, err)
+		req.Header.Set("x-bm-metadata-criticality", string(testCase.crtl))
+		resp, err := client.Do(req)
+		assert.NoError(t, err)
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, testCase.expected, criticalityPkg.Criticality(body))
+	}
+}

pkg/net/metadata/key.go

@@ -36,13 +36,17 @@ const (
 
 	// Device 客户端信息
 	Device = "device"
+
+	// Criticality 重要性
+	Criticality = "criticality"
 )
 
 var outgoingKey = map[string]struct{}{
-	Color:      struct{}{},
-	RemoteIP:   struct{}{},
-	RemotePort: struct{}{},
-	Mirror:     struct{}{},
+	Color:       struct{}{},
+	RemoteIP:    struct{}{},
+	RemotePort:  struct{}{},
+	Mirror:      struct{}{},
+	Criticality: struct{}{},
 }
 
 var incomingKey = map[string]struct{}{