11 Star 11 Fork 0

Gitee 极速下载/goa

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/goadesign/goa
克隆/下载
trace.go 5.44 KB
一键复制 编辑 原始数据 按行查看 历史
Nitin 提交于 2019-01-31 12:54 . [BREAKING] Grpc middleware (#1978)
package middleware
import (
"context"
"goa.design/goa/middleware"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
const (
// TraceIDMetadataKey is the default name of the gRPC request metadata
// key containing the current TraceID if any.
TraceIDMetadataKey = "trace-id"
// ParentSpanIDMetadataKey is the default name of the gRPC request metadata
// key containing the parent span ID if any.
ParentSpanIDMetadataKey = "parent-span-id"
// SpanIDMetadataKey is the default name of the gRPC request metadata
// containing the span ID if any.
SpanIDMetadataKey = "span-id"
)
// UnaryServerTrace returns a server trace middleware that initializes the
// trace informartion in the unary gRPC request context.
//
// Example:
// grpc.NewServer(grpc.UnaryInterceptor(middleware.UnaryServerTrace()))
//
// // enable options
// grpc.NewServer(grpc.UnaryInterceptor(middleware.UnaryServerTrace(
// middleware.TraceIDFunc(myTraceIDFunc),
// middleware.SpanIDFunc(mySpanIDFunc),
// middleware.SamplingPercent(100)))
func UnaryServerTrace(opts ...middleware.TraceOption) grpc.UnaryServerInterceptor {
o := middleware.NewTraceOptions(opts...)
return grpc.UnaryServerInterceptor(func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
ctx = withTrace(ctx, o)
return handler(ctx, req)
})
}
// StreamServerTrace returns a server trace middleware that initializes the
// trace information in the streaming gRPC request context.
//
// Example:
// grpc.NewServer(grpc.StreamInterceptor(middleware.StreamServerTrace()))
//
// // enable options
// grpc.NewServer(grpc.StreamInterceptor(middleware.StreamServerTrace(
// middleware.TraceIDFunc(myTraceIDFunc),
// middleware.SpanIDFunc(mySpanIDFunc),
// middleware.MaxSamplingRate(50)))
func StreamServerTrace(opts ...middleware.TraceOption) grpc.StreamServerInterceptor {
o := middleware.NewTraceOptions(opts...)
return grpc.StreamServerInterceptor(func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
ctx := withTrace(ss.Context(), o)
wss := NewWrappedServerStream(ctx, ss)
return handler(srv, wss)
})
}
// UnaryClientTrace sets the outgoing unary request metadata with the trace
// information found in the context so that the downstream service may properly
// retrieve the parent span ID and trace ID.
//
// Example:
// conn, err := grpc.Dial(url, grpc.WithUnaryInterceptor(UnaryClientTrace()))
func UnaryClientTrace() grpc.UnaryClientInterceptor {
return grpc.UnaryClientInterceptor(func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
ctx = setTrace(ctx)
return invoker(ctx, method, req, reply, cc, opts...)
})
}
// StreamClientTrace sets the outgoing stream request metadata with the trace
// information found in the context so that the downstream service may properly
// retrieve the parent span ID and trace ID.
//
// Example:
// conn, err := grpc.Dial(url, grpc.WithStreamInterceptor(StreamClientTrace()))
func StreamClientTrace() grpc.StreamClientInterceptor {
return grpc.StreamClientInterceptor(func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
ctx = setTrace(ctx)
return streamer(ctx, desc, cc, method, opts...)
})
}
// TraceIDFunc is a wrapper for the top-level TraceIDFunc.
func TraceIDFunc(f middleware.IDFunc) middleware.TraceOption {
return middleware.TraceIDFunc(f)
}
// SpanIDFunc is a wrapper for the top-level SpanIDFunc.
func SpanIDFunc(f middleware.IDFunc) middleware.TraceOption {
return middleware.SpanIDFunc(f)
}
// SamplingPercent is a wrapper for the top-level SamplingPercent.
func SamplingPercent(p int) middleware.TraceOption {
return middleware.SamplingPercent(p)
}
// MaxSamplingRate is a wrapper for the top-level MaxSamplingRate.
func MaxSamplingRate(r int) middleware.TraceOption {
return middleware.MaxSamplingRate(r)
}
// SampleSize is a wrapper for the top-level SampleSize.
func SampleSize(s int) middleware.TraceOption {
return middleware.SampleSize(s)
}
// withTrace sets the trace ID, span ID, and parent span ID in the context.
func withTrace(ctx context.Context, opts *middleware.TraceOptions) context.Context {
sampler := opts.NewSampler()
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
md = metadata.MD{}
}
// insert a new trace ID only if not already being traced.
var traceID string
{
traceID = MetadataValue(md, TraceIDMetadataKey)
if traceID == "" && sampler.Sample() {
// insert tracing only within sample.
traceID = opts.TraceID()
}
}
if traceID == "" {
return ctx
}
var (
spanID string
parentID string
)
{
spanID = opts.SpanID()
parentID = MetadataValue(md, ParentSpanIDMetadataKey)
}
// insert IDs into context to enable tracing.
return middleware.WithSpan(ctx, traceID, spanID, parentID)
}
// setTrace sets the trace information to the request context's outgoing
// metadata.
func setTrace(ctx context.Context) context.Context {
var (
traceID = ctx.Value(middleware.TraceIDKey)
spanID = ctx.Value(middleware.TraceSpanIDKey)
)
if traceID != nil {
md, ok := metadata.FromOutgoingContext(ctx)
if !ok {
md = metadata.MD{}
}
md.Set(TraceIDMetadataKey, traceID.(string))
md.Set(ParentSpanIDMetadataKey, spanID.(string))
ctx = metadata.NewOutgoingContext(ctx, md)
}
return ctx
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/mirrors/goa.git
git@gitee.com:mirrors/goa.git
mirrors
goa
goa
v2.0.0

搜索帮助