github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/pkg/trace/inject/redis/config.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package redis
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"strings"
    21  
    22  	"github.com/go-redis/redis"
    23  	"go.opentelemetry.io/otel"
    24  	"go.opentelemetry.io/otel/attribute"
    25  	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
    26  	"go.opentelemetry.io/otel/trace"
    27  )
    28  
    29  const (
    30  	instrumentationName = "github.com/erda-project/erda-infra/pkg/trace/inject/redis"
    31  )
    32  
    33  // SpanNameFormatter is an interface that used to format span names.
    34  type SpanNameFormatter interface {
    35  	Format(ctx context.Context, cmd redis.Cmder) string
    36  	FormatBatch(ctx context.Context, cmds []redis.Cmder) string
    37  }
    38  
    39  type config struct {
    40  	TracerProvider trace.TracerProvider
    41  	Tracer         trace.Tracer
    42  
    43  	SpanOptions SpanOptions
    44  
    45  	DBSystem string
    46  
    47  	// Attributes will be set to each span.
    48  	Attributes []attribute.KeyValue
    49  
    50  	// SpanNameFormatter will be called to produce span's name.
    51  	// Default use method as span name
    52  	SpanNameFormatter SpanNameFormatter
    53  }
    54  
    55  // SpanOptions holds configuration of tracing span to decide
    56  // whether to enable some features.
    57  // by default all options are set to false intentionally when creating a wrapped
    58  // driver and provide the most sensible default with both performance and
    59  // security in mind.
    60  type SpanOptions struct {
    61  	// Ping, if set to true, will enable the creation of spans on Ping requests.
    62  	Ping bool
    63  
    64  	// DisableStatement if set to true, will suppress db.statement in spans.
    65  	DisableStatement bool
    66  
    67  	// RecordError, if set, will be invoked with the current error, and if the func returns true
    68  	// the record will be recorded on the current span.
    69  	RecordError func(err error) bool
    70  
    71  	// AllowRoot, if set to true, will create root spans in absence of existing spans or even context.
    72  	AllowRoot bool
    73  }
    74  
    75  type defaultSpanNameFormatter struct{}
    76  
    77  func (f *defaultSpanNameFormatter) Format(ctx context.Context, cmd redis.Cmder) string {
    78  	return cmd.Name()
    79  }
    80  
    81  func (f *defaultSpanNameFormatter) FormatBatch(ctx context.Context, cmds []redis.Cmder) string {
    82  	return "batch"
    83  }
    84  
    85  // newConfig returns a config with all Options set.
    86  func newConfig(dbSystem string, options ...Option) *config {
    87  	cfg := config{
    88  		TracerProvider:    otel.GetTracerProvider(),
    89  		DBSystem:          dbSystem,
    90  		SpanNameFormatter: &defaultSpanNameFormatter{},
    91  		SpanOptions:       SpanOptions{},
    92  	}
    93  	for _, opt := range options {
    94  		opt.Apply(&cfg)
    95  	}
    96  
    97  	if cfg.DBSystem != "" {
    98  		cfg.Attributes = append(cfg.Attributes,
    99  			semconv.DBSystemKey.String(cfg.DBSystem),
   100  		)
   101  		cfg.Attributes = cfg.Attributes[:len(cfg.Attributes):len(cfg.Attributes)]
   102  	}
   103  	cfg.Tracer = cfg.TracerProvider.Tracer(
   104  		instrumentationName,
   105  		trace.WithInstrumentationVersion(Version()),
   106  	)
   107  	return &cfg
   108  }
   109  
   110  func withDBStatement(cfg *config, cmd redis.Cmder) []attribute.KeyValue {
   111  	if cfg.SpanOptions.DisableStatement {
   112  		return cfg.Attributes
   113  	}
   114  	return append(cfg.Attributes, semconv.DBStatementKey.String(getStatement(cmd)))
   115  }
   116  
   117  func getStatement(cmd redis.Cmder) string {
   118  	n := len(cmd.Args())
   119  	if n > 1 {
   120  		sb := &strings.Builder{}
   121  		sb.WriteString(cmd.Name())
   122  		sb.WriteString(" ")
   123  		last := n - 2
   124  		for i, arg := range cmd.Args()[1:] {
   125  			sb.WriteString(fmt.Sprint(arg))
   126  			if i < last {
   127  				sb.WriteString(" ")
   128  			}
   129  		}
   130  		return sb.String()
   131  	}
   132  	return cmd.Name()
   133  }
   134  
   135  var statementsCountKey = attribute.Key("db.statements_count")
   136  var isBatchStatementsKey = attribute.Key("db.is_batch_statement")
   137  
   138  func withBatchDBStatement(cfg *config, cmds []redis.Cmder) []attribute.KeyValue {
   139  	if cfg.SpanOptions.DisableStatement {
   140  		return cfg.Attributes
   141  	}
   142  	var statement string
   143  	n := len(cmds)
   144  	switch n {
   145  	case 0:
   146  	case 1:
   147  		statement = getStatement(cmds[0])
   148  	case 2:
   149  		statement = getStatement(cmds[0]) + "; " + getStatement(cmds[1])
   150  	default:
   151  		statement = getStatement(cmds[0]) + "; ... ;" + getStatement(cmds[n-1])
   152  	}
   153  	return append(cfg.Attributes,
   154  		semconv.DBStatementKey.String(statement),
   155  		isBatchStatementsKey.Bool(true),
   156  		statementsCountKey.Int64(int64(n)),
   157  	)
   158  }