github.com/grafana/pyroscope@v1.18.0/pkg/querier/stats/stats.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/querier/stats/stats.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package stats
     7  
     8  import (
     9  	"context"
    10  	"sync/atomic" //lint:ignore faillint we can't use go.uber.org/atomic with a protobuf struct without wrapping it.
    11  	"time"
    12  
    13  	"github.com/grafana/pyroscope/pkg/util/httpgrpc"
    14  )
    15  
    16  type contextKey int
    17  
    18  var ctxKey = contextKey(0)
    19  
    20  // ContextWithEmptyStats returns a context with empty stats.
    21  func ContextWithEmptyStats(ctx context.Context) (*Stats, context.Context) {
    22  	stats := &Stats{}
    23  	ctx = context.WithValue(ctx, ctxKey, stats)
    24  	return stats, ctx
    25  }
    26  
    27  // FromContext gets the Stats out of the Context. Returns nil if stats have not
    28  // been initialised in the context.
    29  func FromContext(ctx context.Context) *Stats {
    30  	o := ctx.Value(ctxKey)
    31  	if o == nil {
    32  		return nil
    33  	}
    34  	return o.(*Stats)
    35  }
    36  
    37  // IsEnabled returns whether stats tracking is enabled in the context.
    38  func IsEnabled(ctx context.Context) bool {
    39  	// When query statistics are enabled, the stats object is already initialised
    40  	// within the context, so we can just check it.
    41  	return FromContext(ctx) != nil
    42  }
    43  
    44  // AddWallTime adds some time to the counter.
    45  func (s *Stats) AddWallTime(t time.Duration) {
    46  	if s == nil {
    47  		return
    48  	}
    49  
    50  	atomic.AddInt64(&s.WallTime, int64(t))
    51  }
    52  
    53  // LoadWallTime returns current wall time.
    54  func (s *Stats) LoadWallTime() time.Duration {
    55  	if s == nil {
    56  		return 0
    57  	}
    58  
    59  	return time.Duration(atomic.LoadInt64((&s.WallTime)))
    60  }
    61  
    62  func (s *Stats) AddFetchedSeries(series uint64) {
    63  	if s == nil {
    64  		return
    65  	}
    66  
    67  	atomic.AddUint64(&s.FetchedSeriesCount, series)
    68  }
    69  
    70  func (s *Stats) LoadFetchedSeries() uint64 {
    71  	if s == nil {
    72  		return 0
    73  	}
    74  
    75  	return atomic.LoadUint64(&s.FetchedSeriesCount)
    76  }
    77  
    78  func (s *Stats) AddFetchedChunkBytes(bytes uint64) {
    79  	if s == nil {
    80  		return
    81  	}
    82  
    83  	atomic.AddUint64(&s.FetchedChunkBytes, bytes)
    84  }
    85  
    86  func (s *Stats) LoadFetchedChunkBytes() uint64 {
    87  	if s == nil {
    88  		return 0
    89  	}
    90  
    91  	return atomic.LoadUint64(&s.FetchedChunkBytes)
    92  }
    93  
    94  func (s *Stats) AddFetchedChunks(chunks uint64) {
    95  	if s == nil {
    96  		return
    97  	}
    98  
    99  	atomic.AddUint64(&s.FetchedChunksCount, chunks)
   100  }
   101  
   102  func (s *Stats) LoadFetchedChunks() uint64 {
   103  	if s == nil {
   104  		return 0
   105  	}
   106  
   107  	return atomic.LoadUint64(&s.FetchedChunksCount)
   108  }
   109  
   110  func (s *Stats) AddFetchedIndexBytes(indexBytes uint64) {
   111  	if s == nil {
   112  		return
   113  	}
   114  
   115  	atomic.AddUint64(&s.FetchedIndexBytes, indexBytes)
   116  }
   117  
   118  func (s *Stats) LoadFetchedIndexBytes() uint64 {
   119  	if s == nil {
   120  		return 0
   121  	}
   122  
   123  	return atomic.LoadUint64(&s.FetchedIndexBytes)
   124  }
   125  
   126  func (s *Stats) AddShardedQueries(num uint32) {
   127  	if s == nil {
   128  		return
   129  	}
   130  
   131  	atomic.AddUint32(&s.ShardedQueries, num)
   132  }
   133  
   134  func (s *Stats) LoadShardedQueries() uint32 {
   135  	if s == nil {
   136  		return 0
   137  	}
   138  
   139  	return atomic.LoadUint32(&s.ShardedQueries)
   140  }
   141  
   142  func (s *Stats) AddSplitQueries(num uint32) {
   143  	if s == nil {
   144  		return
   145  	}
   146  
   147  	atomic.AddUint32(&s.SplitQueries, num)
   148  }
   149  
   150  func (s *Stats) LoadSplitQueries() uint32 {
   151  	if s == nil {
   152  		return 0
   153  	}
   154  
   155  	return atomic.LoadUint32(&s.SplitQueries)
   156  }
   157  
   158  // Merge the provided Stats into this one.
   159  func (s *Stats) Merge(other *Stats) {
   160  	if s == nil || other == nil {
   161  		return
   162  	}
   163  
   164  	s.AddWallTime(other.LoadWallTime())
   165  	s.AddFetchedSeries(other.LoadFetchedSeries())
   166  	s.AddFetchedChunkBytes(other.LoadFetchedChunkBytes())
   167  	s.AddFetchedChunks(other.LoadFetchedChunks())
   168  	s.AddShardedQueries(other.LoadShardedQueries())
   169  	s.AddSplitQueries(other.LoadSplitQueries())
   170  	s.AddFetchedIndexBytes(other.LoadFetchedIndexBytes())
   171  }
   172  
   173  func ShouldTrackHTTPGRPCResponse(r *httpgrpc.HTTPResponse) bool {
   174  	// Do no track statistics for requests failed because of a server error.
   175  	return r.Code < 500
   176  }