vitess.io/vitess@v0.16.2/go/trace/opentracing.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package trace
    18  
    19  import (
    20  	"context"
    21  	"encoding/base64"
    22  	"encoding/json"
    23  
    24  	otgrpc "github.com/opentracing-contrib/go-grpc"
    25  	"github.com/opentracing/opentracing-go"
    26  	"google.golang.org/grpc"
    27  
    28  	"vitess.io/vitess/go/vt/vterrors"
    29  )
    30  
    31  var _ Span = (*openTracingSpan)(nil)
    32  
    33  type openTracingSpan struct {
    34  	otSpan opentracing.Span
    35  }
    36  
    37  // Finish will mark a span as finished
    38  func (js openTracingSpan) Finish() {
    39  	js.otSpan.Finish()
    40  }
    41  
    42  // Annotate will add information to an existing span
    43  func (js openTracingSpan) Annotate(key string, value any) {
    44  	js.otSpan.SetTag(key, value)
    45  }
    46  
    47  var _ tracingService = (*openTracingService)(nil)
    48  
    49  type tracer interface {
    50  	GetOpenTracingTracer() opentracing.Tracer
    51  }
    52  
    53  type openTracingService struct {
    54  	Tracer tracer
    55  }
    56  
    57  // AddGrpcServerOptions is part of an interface implementation
    58  func (jf openTracingService) AddGrpcServerOptions(addInterceptors func(s grpc.StreamServerInterceptor, u grpc.UnaryServerInterceptor)) {
    59  	ot := jf.Tracer.GetOpenTracingTracer()
    60  	addInterceptors(otgrpc.OpenTracingStreamServerInterceptor(ot), otgrpc.OpenTracingServerInterceptor(ot))
    61  }
    62  
    63  // AddGrpcClientOptions is part of an interface implementation
    64  func (jf openTracingService) AddGrpcClientOptions(addInterceptors func(s grpc.StreamClientInterceptor, u grpc.UnaryClientInterceptor)) {
    65  	ot := jf.Tracer.GetOpenTracingTracer()
    66  	addInterceptors(otgrpc.OpenTracingStreamClientInterceptor(ot), otgrpc.OpenTracingClientInterceptor(ot))
    67  }
    68  
    69  // NewClientSpan is part of an interface implementation
    70  func (jf openTracingService) NewClientSpan(parent Span, serviceName, label string) Span {
    71  	span := jf.New(parent, label)
    72  	span.Annotate("peer.service", serviceName)
    73  	return span
    74  }
    75  
    76  // New is part of an interface implementation
    77  func (jf openTracingService) New(parent Span, label string) Span {
    78  	var innerSpan opentracing.Span
    79  	if parent == nil {
    80  		innerSpan = jf.Tracer.GetOpenTracingTracer().StartSpan(label)
    81  	} else {
    82  		jaegerParent := parent.(openTracingSpan)
    83  		span := jaegerParent.otSpan
    84  		innerSpan = jf.Tracer.GetOpenTracingTracer().StartSpan(label, opentracing.ChildOf(span.Context()))
    85  	}
    86  	return openTracingSpan{otSpan: innerSpan}
    87  }
    88  
    89  func extractMapFromString(in string) (opentracing.TextMapCarrier, error) {
    90  	decodedBytes, err := base64.StdEncoding.DecodeString(in)
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  
    95  	var dat opentracing.TextMapCarrier
    96  	err = json.Unmarshal(decodedBytes, &dat)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	return dat, nil
   102  }
   103  
   104  func (jf openTracingService) NewFromString(parent, label string) (Span, error) {
   105  	carrier, err := extractMapFromString(parent)
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  	spanContext, err := jf.Tracer.GetOpenTracingTracer().Extract(opentracing.TextMap, carrier)
   110  	if err != nil {
   111  		return nil, vterrors.Wrap(err, "failed to deserialize span context")
   112  	}
   113  	innerSpan := jf.Tracer.GetOpenTracingTracer().StartSpan(label, opentracing.ChildOf(spanContext))
   114  	return openTracingSpan{otSpan: innerSpan}, nil
   115  }
   116  
   117  // FromContext is part of an interface implementation
   118  func (jf openTracingService) FromContext(ctx context.Context) (Span, bool) {
   119  	innerSpan := opentracing.SpanFromContext(ctx)
   120  
   121  	if innerSpan == nil {
   122  		return nil, false
   123  	}
   124  	return openTracingSpan{otSpan: innerSpan}, true
   125  }
   126  
   127  // NewContext is part of an interface implementation
   128  func (jf openTracingService) NewContext(parent context.Context, s Span) context.Context {
   129  	span, ok := s.(openTracingSpan)
   130  	if !ok {
   131  		return nil
   132  	}
   133  	return opentracing.ContextWithSpan(parent, span.otSpan)
   134  }