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

/**
 * @Author: mac
 * @Description: grpc opentracing clientInterceptor
 * @Date: 2020/3/25 10:07
 */

package trace

import (
	"context"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/opentracing/opentracing-go/log"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func ClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string,
		req, reply interface{}, cc *grpc.ClientConn,
		invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {

		span, ctx := opentracing.StartSpanFromContext(
			ctx,
			"GRPC client ",
			opentracing.Tag{Key: string(ext.Component), Value: "GRPC"},
			ext.SpanKindRPCClient,
		)

		defer span.Finish()

		md, ok := metadata.FromOutgoingContext(ctx)
		if !ok {
			md = metadata.New(nil)
		} else {
			md = md.Copy()
		}

		err := opentracing.GlobalTracer().Inject(span.Context(), opentracing.TextMap, MDReaderWriter{md})
		if err != nil {
			span.SetTag("error", true)
			span.LogFields(log.String("inject-error", err.Error()))
		}

		newCtx := metadata.NewOutgoingContext(ctx, md)
		err = invoker(newCtx, method, req, reply, cc, opts...)
		if err != nil {
			span.SetTag("error", true)
			span.LogFields(
				log.Object("req", req),
				log.Object("reply", reply),
				log.String("call-error", err.Error()))
		}
		return err
	}
}