57 lines
1.4 KiB
Go
57 lines
1.4 KiB
Go
// Copyright (C) 2019 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
// Package rpctracing implements tracing for rpc.
|
|
package rpctracing
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/spacemonkeygo/monkit/v3"
|
|
|
|
"storj.io/drpc"
|
|
"storj.io/drpc/drpcmetadata"
|
|
"storj.io/drpc/drpcmux"
|
|
)
|
|
|
|
type streamWrapper struct {
|
|
drpc.Stream
|
|
ctx context.Context
|
|
}
|
|
|
|
func (s *streamWrapper) Context() context.Context { return s.ctx }
|
|
|
|
type handlerFunc func(metadata map[string]string) (trace *monkit.Trace, spanID int64)
|
|
|
|
func defaultHandlerFunc(metadata map[string]string) (*monkit.Trace, int64) {
|
|
return monkit.NewTrace(monkit.NewId()), monkit.NewId()
|
|
}
|
|
|
|
// Handler implements drpc handler interface and takes in a callback function.
|
|
type Handler struct {
|
|
mux *drpcmux.Mux
|
|
cb handlerFunc
|
|
}
|
|
|
|
// NewHandler returns a new instance of Handler.
|
|
func NewHandler(mux *drpcmux.Mux, cb handlerFunc) *Handler {
|
|
if cb == nil {
|
|
cb = defaultHandlerFunc
|
|
}
|
|
return &Handler{
|
|
mux: mux,
|
|
cb: cb,
|
|
}
|
|
}
|
|
|
|
// HandleRPC adds tracing metadata onto server stream.
|
|
func (handler *Handler) HandleRPC(stream drpc.Stream, rpc string) (err error) {
|
|
streamCtx := stream.Context()
|
|
metadata, ok := drpcmetadata.Get(streamCtx)
|
|
if ok {
|
|
trace, spanID := handler.cb(metadata)
|
|
defer mon.FuncNamed(rpc).RemoteTrace(&streamCtx, spanID, trace)(&err)
|
|
}
|
|
|
|
return handler.mux.HandleRPC(&streamWrapper{Stream: stream, ctx: streamCtx}, rpc)
|
|
}
|