github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/stats/stats.go (about)

     1  package stats
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic" //lint:ignore faillint we can't use go.uber.org/atomic with a protobuf struct without wrapping it.
     6  	"time"
     7  
     8  	"github.com/weaveworks/common/httpgrpc"
     9  )
    10  
    11  type contextKey int
    12  
    13  var ctxKey = contextKey(0)
    14  
    15  // ContextWithEmptyStats returns a context with empty stats.
    16  func ContextWithEmptyStats(ctx context.Context) (*Stats, context.Context) {
    17  	stats := &Stats{}
    18  	ctx = context.WithValue(ctx, ctxKey, stats)
    19  	return stats, ctx
    20  }
    21  
    22  // FromContext gets the Stats out of the Context. Returns nil if stats have not
    23  // been initialised in the context.
    24  func FromContext(ctx context.Context) *Stats {
    25  	o := ctx.Value(ctxKey)
    26  	if o == nil {
    27  		return nil
    28  	}
    29  	return o.(*Stats)
    30  }
    31  
    32  // IsEnabled returns whether stats tracking is enabled in the context.
    33  func IsEnabled(ctx context.Context) bool {
    34  	// When query statistics are enabled, the stats object is already initialised
    35  	// within the context, so we can just check it.
    36  	return FromContext(ctx) != nil
    37  }
    38  
    39  // AddWallTime adds some time to the counter.
    40  func (s *Stats) AddWallTime(t time.Duration) {
    41  	if s == nil {
    42  		return
    43  	}
    44  
    45  	atomic.AddInt64((*int64)(&s.WallTime), int64(t))
    46  }
    47  
    48  // LoadWallTime returns current wall time.
    49  func (s *Stats) LoadWallTime() time.Duration {
    50  	if s == nil {
    51  		return 0
    52  	}
    53  
    54  	return time.Duration(atomic.LoadInt64((*int64)(&s.WallTime)))
    55  }
    56  
    57  func (s *Stats) AddFetchedSeries(series uint64) {
    58  	if s == nil {
    59  		return
    60  	}
    61  
    62  	atomic.AddUint64(&s.FetchedSeriesCount, series)
    63  }
    64  
    65  func (s *Stats) LoadFetchedSeries() uint64 {
    66  	if s == nil {
    67  		return 0
    68  	}
    69  
    70  	return atomic.LoadUint64(&s.FetchedSeriesCount)
    71  }
    72  
    73  func (s *Stats) AddFetchedChunkBytes(bytes uint64) {
    74  	if s == nil {
    75  		return
    76  	}
    77  
    78  	atomic.AddUint64(&s.FetchedChunkBytes, bytes)
    79  }
    80  
    81  func (s *Stats) LoadFetchedChunkBytes() uint64 {
    82  	if s == nil {
    83  		return 0
    84  	}
    85  
    86  	return atomic.LoadUint64(&s.FetchedChunkBytes)
    87  }
    88  
    89  // Merge the provide Stats into this one.
    90  func (s *Stats) Merge(other *Stats) {
    91  	if s == nil || other == nil {
    92  		return
    93  	}
    94  
    95  	s.AddWallTime(other.LoadWallTime())
    96  	s.AddFetchedSeries(other.LoadFetchedSeries())
    97  	s.AddFetchedChunkBytes(other.LoadFetchedChunkBytes())
    98  }
    99  
   100  func ShouldTrackHTTPGRPCResponse(r *httpgrpc.HTTPResponse) bool {
   101  	// Do no track statistics for requests failed because of a server error.
   102  	return r.Code < 500
   103  }