github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/logger/zaputil/fieldsampler.go (about)

     1  // Copyright 2023 LiveKit, 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 zaputil
    16  
    17  import (
    18  	"fmt"
    19  	"math"
    20  	"reflect"
    21  	"sync/atomic"
    22  
    23  	"github.com/zeebo/xxh3"
    24  	"go.uber.org/zap"
    25  	"go.uber.org/zap/zapcore"
    26  )
    27  
    28  type FieldSampleRate interface {
    29  	Threshold() uint64
    30  }
    31  
    32  func rateToThreshold(rate float64) uint64 {
    33  	return math.MaxUint64 / 10000 * uint64(math.Max(math.Min(rate, 1), 0)*10000)
    34  }
    35  
    36  type AtomicFieldSampleRate uint64
    37  
    38  func NewAtomicFieldSampleRate(rate float64) *AtomicFieldSampleRate {
    39  	var r AtomicFieldSampleRate
    40  	r.SetRate(rate)
    41  	return &r
    42  }
    43  
    44  func (r *AtomicFieldSampleRate) SetRate(rate float64) {
    45  	atomic.StoreUint64((*uint64)(r), rateToThreshold(rate))
    46  }
    47  
    48  func (r *AtomicFieldSampleRate) Threshold() uint64 {
    49  	return uint64(atomic.LoadUint64((*uint64)(r)))
    50  }
    51  
    52  type FieldSamplerAction int
    53  
    54  const (
    55  	OmitSampledLog FieldSamplerAction = iota
    56  	AnnotateSampledLog
    57  )
    58  
    59  type FieldSamplerConfig struct {
    60  	FieldName           string
    61  	Rate                FieldSampleRate
    62  	Action              FieldSamplerAction
    63  	AnnotationFieldName string
    64  }
    65  
    66  func NewFieldSampler(core zapcore.Core, config FieldSamplerConfig) zapcore.Core {
    67  	return &fieldSampler{
    68  		Core:   core,
    69  		config: config,
    70  	}
    71  }
    72  
    73  type fieldSampler struct {
    74  	zapcore.Core
    75  	config FieldSamplerConfig
    76  	hash   uint64
    77  }
    78  
    79  var _ zapcore.Core = (*fieldSampler)(nil)
    80  
    81  func (s *fieldSampler) With(fields []zapcore.Field) zapcore.Core {
    82  	hash := s.hash
    83  	if h, ok := s.hashSampleField(fields); ok {
    84  		hash = h
    85  	}
    86  
    87  	return &fieldSampler{
    88  		Core:   s.Core.With(fields),
    89  		config: s.config,
    90  		hash:   hash,
    91  	}
    92  }
    93  
    94  func (s *fieldSampler) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
    95  	if s.Enabled(ent.Level) {
    96  		return ce.AddCore(ent, s)
    97  	}
    98  	return ce
    99  }
   100  
   101  func (s *fieldSampler) hashSampleField(fields []zapcore.Field) (uint64, bool) {
   102  	v, ok := s.findSampleField(fields)
   103  	if !ok {
   104  		return 0, false
   105  	}
   106  	return xxh3.HashString(v), true
   107  }
   108  
   109  func (s *fieldSampler) findSampleField(fields []zapcore.Field) (string, bool) {
   110  	for _, f := range fields {
   111  		if f.Key == s.config.FieldName {
   112  			switch f.Type {
   113  			case zapcore.StringType:
   114  				return f.String, true
   115  			case zapcore.ReflectType:
   116  				rv := reflect.ValueOf(f.Interface)
   117  				if rv.Kind() == reflect.String {
   118  					return rv.String(), true
   119  				}
   120  			case zapcore.StringerType:
   121  				if str, ok := f.Interface.(fmt.Stringer); ok {
   122  					return str.String(), true
   123  				}
   124  			}
   125  			return "", false
   126  		}
   127  	}
   128  	return "", false
   129  }
   130  
   131  func (s *fieldSampler) test(fields []zapcore.Field) bool {
   132  	if s.hash != 0 {
   133  		return s.hash > s.config.Rate.Threshold()
   134  	}
   135  	if h, ok := s.hashSampleField(fields); ok {
   136  		return h > s.config.Rate.Threshold()
   137  	}
   138  	return false
   139  }
   140  
   141  func (s *fieldSampler) Write(entry zapcore.Entry, fields []zapcore.Field) error {
   142  	if s.test(fields) {
   143  		switch s.config.Action {
   144  		case OmitSampledLog:
   145  			return nil
   146  		case AnnotateSampledLog:
   147  			fields = append(fields, zap.Bool(s.config.AnnotationFieldName, true))
   148  		}
   149  	}
   150  	return s.Core.Write(entry, fields)
   151  }