github.com/thanos-io/thanos@v0.32.5/internal/cortex/querier/stats/stats.go (about)

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