Skip to content

Commit

Permalink
chore: rename (#17)
Browse files Browse the repository at this point in the history
* chore: rename

* chore: rename
  • Loading branch information
chenquan authored Mar 4, 2023
1 parent 69d94ee commit e2c0609
Show file tree
Hide file tree
Showing 7 changed files with 73 additions and 67 deletions.
57 changes: 57 additions & 0 deletions internal/tag/tag.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
package tag

import (
"context"

"github.com/zeromicro/go-zero/core/logx"
"go.opentelemetry.io/otel/baggage"
"google.golang.org/grpc/attributes"
)

const Key = "x-zero-flow-tag"

func ContextWithTag(ctx context.Context, tag string) context.Context {
bg := baggage.FromContext(ctx)
member, err := baggage.NewMember(Key, tag)
if err != nil {
logx.WithContext(ctx).Error(err)
return ctx
}

bg, err = bg.SetMember(member)
if err != nil {
logx.WithContext(ctx).Error(err)
return ctx
}

ctx = baggage.ContextWithBaggage(ctx, bg)

return ctx
}

func FromContext(ctx context.Context) string {
bg := baggage.FromContext(ctx)
member := bg.Member(Key)

return member.Value()
}

type tagAttributesKey struct{}

func NewAttributes(tag string) *attributes.Attributes {
return attributes.New(tagAttributesKey{}, tag)
}

func FromGrpcAttributes(attributes *attributes.Attributes) (string, bool) {
value := attributes.Value(tagAttributesKey{})
if value == nil {
return "", false
}

m, ok := value.(string)
if !ok {
return "", false
}

return m, true
}
12 changes: 6 additions & 6 deletions rest/internal/handler/tag_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,31 +4,31 @@ import (
"context"
"net/http"

"github.com/chenquan/zero-flow/tag"
"github.com/chenquan/zero-flow/internal/tag"
"github.com/zeromicro/go-zero/core/logx"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
)

var httpColorAttributeKey = attribute.Key("http.header.color")

func TagHandler(headerTag string) func(next http.Handler) http.Handler {
func TagHandler(tagHeader string) func(next http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
ctx := request.Context()
ctx = newBaggage(ctx, request, headerTag)
ctx = newBaggage(ctx, request, tagHeader)
next.ServeHTTP(writer, request.WithContext(ctx))
})
}
}

func newBaggage(ctx context.Context, request *http.Request, headerTag string) context.Context {
func newBaggage(ctx context.Context, request *http.Request, tagHeader string) context.Context {
span := trace.SpanFromContext(ctx)
tagString := request.Header.Get(headerTag)
tagString := request.Header.Get(tagHeader)
if len(tagString) == 0 {
return ctx
}
logx.WithContext(ctx).Debugw("flow staining...", logx.Field("tag", tagString))
logx.WithContext(ctx).Debugw("flow staining...", logx.Field(tag.Key, tagString))

ctx = tag.ContextWithTag(ctx, tagString)
span.SetAttributes(httpColorAttributeKey.String(tagString))
Expand Down
4 changes: 2 additions & 2 deletions rest/rest.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@ type (
Server = rest.Server
RestConf struct {
rest.RestConf
HeaderTag string `json:",optional"`
TagHeader string `json:",optional"`
}
)

func MustNewServer(c RestConf, opts ...RunOption) *Server {
server := rest.MustNewServer(c.RestConf, opts...)
server.Use(rest.ToMiddleware(handler.TagHandler(c.HeaderTag)))
server.Use(rest.ToMiddleware(handler.TagHandler(c.TagHeader)))
return server
}
57 changes: 4 additions & 53 deletions tag/tag.go
Original file line number Diff line number Diff line change
@@ -1,57 +1,8 @@
package tag

import (
"context"
import "github.com/chenquan/zero-flow/internal/tag"

"github.com/zeromicro/go-zero/core/logx"
"go.opentelemetry.io/otel/baggage"
"google.golang.org/grpc/attributes"
var (
ContextWithTag = tag.ContextWithTag
FromContext = tag.FromContext
)

const tagKey = "x-zero-flow-tag"

func ContextWithTag(ctx context.Context, tag string) context.Context {
bg := baggage.FromContext(ctx)
member, err := baggage.NewMember(tagKey, tag)
if err != nil {
logx.WithContext(ctx).Error(err)
return ctx
}

bg, err = bg.SetMember(member)
if err != nil {
logx.WithContext(ctx).Error(err)
return ctx
}

ctx = baggage.ContextWithBaggage(ctx, bg)

return ctx
}

func FromContext(ctx context.Context) string {
bg := baggage.FromContext(ctx)
member := bg.Member(tagKey)

return member.Value()
}

type tagAttributesKey struct{}

func NewAttributes(tag string) *attributes.Attributes {
return attributes.New(tagAttributesKey{}, tag)
}

func FromGrpcAttributes(attributes *attributes.Attributes) (string, bool) {
value := attributes.Value(tagAttributesKey{})
if value == nil {
return "", false
}

m, ok := value.(string)
if !ok {
return "", false
}

return m, true
}
2 changes: 1 addition & 1 deletion zrpc/internal/p2c/p2c.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
"sync/atomic"
"time"

"github.com/chenquan/zero-flow/tag"
"github.com/chenquan/zero-flow/internal/tag"
"github.com/chenquan/zero-flow/zrpc/internal/selector"
"github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/core/syncx"
Expand Down
2 changes: 1 addition & 1 deletion zrpc/internal/resolver/discovbuilder.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import (
"net/url"
"strings"

"github.com/chenquan/zero-flow/tag"
"github.com/chenquan/zero-flow/internal/tag"
"github.com/chenquan/zero-flow/zrpc/internal/targets"
"github.com/zeromicro/go-zero/core/discov"
"github.com/zeromicro/go-zero/core/logx"
Expand Down
6 changes: 2 additions & 4 deletions zrpc/internal/selector/defaultselector.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,12 @@
package selector

import (
"github.com/chenquan/zero-flow/tag"
"github.com/chenquan/zero-flow/internal/tag"
"github.com/zeromicro/go-zero/core/logx"
"go.opentelemetry.io/otel/trace"
"google.golang.org/grpc/balancer"
)

const tagKey = "tag"

var (
DefaultSelector = defaultSelector{}
_ Selector = (*defaultSelector)(nil)
Expand All @@ -30,7 +28,7 @@ func (d defaultSelector) Select(conns []Conn, info balancer.PickInfo) []Conn {
}

if len(newConns) != 0 {
logx.WithContext(info.Ctx).Debugw("flow staining...", logx.Field(tagKey, tagString))
logx.WithContext(info.Ctx).Debugw("flow staining...", logx.Field(tag.Key, tagString))

spanCtx := trace.SpanFromContext(info.Ctx)
spanCtx.SetAttributes(tagAttributeKey.String(tagString))
Expand Down

0 comments on commit e2c0609

Please sign in to comment.