58 lines
1.5 KiB
Go
58 lines
1.5 KiB
Go
// Copyright (C) 2019 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
// Package rpctracing implements tracing for rpc.
|
|
package rpctracing
|
|
|
|
import (
|
|
"context"
|
|
"strconv"
|
|
|
|
"github.com/spacemonkeygo/monkit/v3"
|
|
|
|
"storj.io/drpc"
|
|
"storj.io/drpc/drpcmetadata"
|
|
)
|
|
|
|
// TracingWrapper wraps a drpc.Conn with tracing information.
|
|
type TracingWrapper struct {
|
|
drpc.Conn
|
|
}
|
|
|
|
// NewTracingWrapper creates a new instance of the wrapper.
|
|
func NewTracingWrapper(conn drpc.Conn) *TracingWrapper {
|
|
return &TracingWrapper{
|
|
conn,
|
|
}
|
|
}
|
|
|
|
// Invoke implements drpc.Conn's Invoke method with tracing information injected into the context.
|
|
func (c *TracingWrapper) Invoke(ctx context.Context, rpc string, in drpc.Message, out drpc.Message) (err error) {
|
|
return c.Conn.Invoke(c.trace(ctx), rpc, in, out)
|
|
}
|
|
|
|
// NewStream implements drpc.Conn's NewStream method with tracing information injected into the context.
|
|
func (c *TracingWrapper) NewStream(ctx context.Context, rpc string) (_ drpc.Stream, err error) {
|
|
return c.Conn.NewStream(c.trace(ctx), rpc)
|
|
}
|
|
|
|
// trace injects tracing related information into the context.
|
|
func (c *TracingWrapper) trace(ctx context.Context) context.Context {
|
|
span := monkit.SpanFromCtx(ctx)
|
|
if span == nil || span.Parent() == nil {
|
|
return ctx
|
|
}
|
|
|
|
sampled, exist := span.Trace().Get(Sampled).(bool)
|
|
if !exist || !sampled {
|
|
return ctx
|
|
}
|
|
|
|
data := map[string]string{
|
|
TraceID: strconv.FormatInt(span.Trace().Id(), 10),
|
|
ParentID: strconv.FormatInt(span.Id(), 10),
|
|
Sampled: strconv.FormatBool(sampled),
|
|
}
|
|
|
|
return drpcmetadata.AddPairs(ctx, data)
|
|
}
|