gopkgsv2/trace/rpcx_client.go

// Copyright 2020 getensh.com. All rights reserved.
// Use of this source code is governed by getensh.com.

/**
 * @Author: mac
 * @Description: rpcx opentracing clientInterceptor
 * @Date: 2020/3/25 10:20
 */

package trace

import (
	"context"
	"github.com/opentracing/opentracing-go/log"

	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/smallnest/rpcx/share"
)

type RpcxOpenTracingClientPlugin struct{}

func (p *RpcxOpenTracingClientPlugin) DoPreCall(ctx context.Context, servicePath, serviceMethod string, args interface{}) error {
	var span1 opentracing.Span

	// if it is called in rpc service in case that a service calls antoher service,
	// we uses the span in the service context as the parent span.
	parentSpan := ctx.Value(share.OpentracingSpanServerKey)
	if parentSpan != nil {
		span1 = opentracing.StartSpan(
			"rpcx.client."+servicePath+"."+serviceMethod,
			opentracing.ChildOf(parentSpan.(opentracing.Span).Context()))
	} else {
		wireContext, err := share.GetSpanContextFromContext(ctx)
		if err == nil && wireContext != nil { //try to parse span from request
			span1 = opentracing.StartSpan(
				"rpcx.client."+servicePath+"."+serviceMethod,
				ext.RPCServerOption(wireContext))
		} else { // parse span from context or create root context
			span1, _ = opentracing.StartSpanFromContext(ctx, "rpcx.client."+servicePath+"."+serviceMethod)
		}
	}

	if rpcxContext, ok := ctx.(*share.Context); ok {
		rpcxContext.SetValue(share.OpentracingSpanClientKey, span1)
	}
	return nil
}

func (p *RpcxOpenTracingClientPlugin) DoPostCall(ctx context.Context, servicePath, serviceMethod string, args interface{}, reply interface{}, err error) error {
	if rpcxContext, ok := ctx.(*share.Context); ok {
		span1 := rpcxContext.Value(share.OpentracingSpanClientKey)
		if span1 != nil {
			sp := span1.(opentracing.Span)
			if err != nil {
				sp.SetTag("error", true)
				sp.LogFields(
					log.Object("request", args),
					log.Object("reply", reply),
					log.Error(err))
			}
			sp.Finish()
		}
	}
	return nil
}