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

/**
 * @Author: mac
 * @Description: grpc opentracing serverInterceptor
 * @Date: 2020/3/25 10:08
 */

package trace

import (
	"context"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/uber/jaeger-client-go"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/metadata"
)

// grpc服务端启动选项
func ServerOption() grpc.ServerOption {
	return grpc.UnaryInterceptor(serverInterceptor(opentracing.GlobalTracer()))
}

// grpc服务端拦截器
func serverInterceptor(tracer opentracing.Tracer) grpc.UnaryServerInterceptor {
	return func(ctx context.Context,
		req interface{},
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp interface{}, err error) {

		var parentContext context.Context

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

		spanContext, err := tracer.Extract(opentracing.TextMap, MDReaderWriter{md})
		if err != nil && err != opentracing.ErrSpanContextNotFound {
			grpclog.Errorf("extract from metadata err: %v", err)
		} else {
			span := tracer.StartSpan(
				"GRPC server "+info.FullMethod,
				ext.RPCServerOption(spanContext),
				opentracing.Tag{Key: string(ext.Component), Value: "GRPC"},
				ext.SpanKindRPCServer,
			)
			defer span.Finish()

			parentContext = opentracing.ContextWithSpan(ctx, span)

			// 设置 request ID,可用于日志记录, 使用ctx.Value("RequestId")可以取出
			if sc, ok := span.Context().(jaeger.SpanContext); ok {
				traceId := sc.TraceID().String()
				parentContext = context.WithValue(parentContext, "RequestId", traceId)
			}
		}

		return handler(parentContext, req)
	}
}