github.com/tsuna/gohbase@v0.0.0-20250731002811-4ffcadfba63e/internal/observability/observability.go (about)

     1  // Copyright (C) 2021  The GoHBase Authors.  All rights reserved.
     2  // This file is part of GoHBase.
     3  // Use of this source code is governed by the Apache License 2.0
     4  // that can be found in the COPYING file.
     5  
     6  package observability
     7  
     8  import (
     9  	"context"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/tsuna/gohbase/pb"
    13  	"go.opentelemetry.io/otel"
    14  	"go.opentelemetry.io/otel/trace"
    15  )
    16  
    17  // RequestTracePropagator is used to propagate
    18  // tracing information into the RPCTInfo field
    19  // of the RequestHeader
    20  type RequestTracePropagator struct {
    21  	RequestHeader *pb.RequestHeader
    22  }
    23  
    24  // StartSpan starts a trace with the given context
    25  func StartSpan(
    26  	ctx context.Context,
    27  	name string,
    28  	opts ...trace.SpanStartOption,
    29  ) (context.Context, trace.Span) {
    30  	return otel.Tracer("gohbase").Start(ctx, name, opts...)
    31  }
    32  
    33  // ObserveWithTrace observes the value, providing the traceID as
    34  // an exemplar if exemplars are supported and a traceID is present
    35  func ObserveWithTrace(ctx context.Context, o prometheus.Observer, v float64) {
    36  	spanContext := trace.SpanContextFromContext(ctx)
    37  	if spanContext.IsSampled() && spanContext.HasTraceID() {
    38  		traceID := spanContext.TraceID().String()
    39  
    40  		if exemplarObserver, ok := o.(prometheus.ExemplarObserver); ok {
    41  			exemplarObserver.ObserveWithExemplar(v, prometheus.Labels{
    42  				"traceID": traceID,
    43  			})
    44  			return
    45  		}
    46  	}
    47  
    48  	o.Observe(v)
    49  }
    50  
    51  // Get implements the go.opentelemetry.io/otel/propagation.TextMapCarrier interface
    52  func (r RequestTracePropagator) Get(key string) string {
    53  	if r.RequestHeader == nil ||
    54  		r.RequestHeader.TraceInfo == nil ||
    55  		r.RequestHeader.TraceInfo.Headers == nil {
    56  		return ""
    57  	}
    58  	return r.RequestHeader.TraceInfo.Headers[key]
    59  }
    60  
    61  // Set implements the go.opentelemetry.io/otel/propagation.TextMapCarrier interface
    62  func (r RequestTracePropagator) Set(key string, value string) {
    63  	if r.RequestHeader == nil {
    64  		return
    65  	}
    66  	if r.RequestHeader.TraceInfo == nil {
    67  		r.RequestHeader.TraceInfo = &pb.RPCTInfo{
    68  			Headers: make(map[string]string),
    69  		}
    70  	}
    71  	r.RequestHeader.TraceInfo.Headers[key] = value
    72  }
    73  
    74  // Keys implements the go.opentelemetry.io/otel/propagation.TextMapCarrier interface
    75  func (r RequestTracePropagator) Keys() []string {
    76  	if r.RequestHeader == nil ||
    77  		r.RequestHeader.TraceInfo == nil ||
    78  		r.RequestHeader.TraceInfo.Headers == nil {
    79  		return []string{}
    80  	}
    81  
    82  	h := r.RequestHeader.TraceInfo.Headers
    83  	result := make([]string, 0, len(h))
    84  	for k := range h {
    85  		result = append(result, k)
    86  	}
    87  	return result
    88  }