github.com/keybase/client/go@v0.0.0-20241007131713-f10651d043c8/profiling/timebuckets.go (about)

     1  package profiling
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/keybase/client/go/logger"
     8  	"github.com/keybase/clockwork"
     9  	"golang.org/x/net/context"
    10  )
    11  
    12  type ctxKeyType string
    13  
    14  var ctxKey = ctxKeyType("timebuckets")
    15  
    16  type TimeBuckets struct {
    17  	sync.Mutex
    18  	clock clockwork.Clock
    19  	log   logger.Logger
    20  	times map[string]time.Duration
    21  }
    22  
    23  func NewTimeBuckets(clock clockwork.Clock, log logger.Logger) *TimeBuckets {
    24  	return &TimeBuckets{
    25  		clock: clock,
    26  		log:   log,
    27  		times: make(map[string]time.Duration),
    28  	}
    29  }
    30  
    31  func (t *TimeBuckets) Record(bucketName string) FinFn {
    32  	start := t.clock.Now()
    33  	return func() {
    34  		duration := t.clock.Since(start)
    35  		t.Lock()
    36  		defer t.Unlock()
    37  		t.times[bucketName] += duration
    38  	}
    39  }
    40  
    41  func (t *TimeBuckets) Get(bucketName string) time.Duration {
    42  	t.Lock()
    43  	defer t.Unlock()
    44  	return t.times[bucketName]
    45  }
    46  
    47  func (t *TimeBuckets) Log(ctx context.Context, bucketName string) {
    48  	t.log.CDebugf(ctx, "TimeBucket %s [time=%s]", bucketName, t.Get(bucketName))
    49  }
    50  
    51  func (t *TimeBuckets) LogIfNonZero(ctx context.Context, bucketName string) {
    52  	d := t.Get(bucketName)
    53  	if d != 0 {
    54  		t.log.CDebugf(ctx, "TimeBucket %s [time=%s]", bucketName, d)
    55  	}
    56  }
    57  
    58  type FinFn func()
    59  
    60  func WithTimeBuckets(ctx context.Context, clock clockwork.Clock, log logger.Logger) (context.Context, *TimeBuckets) {
    61  	v, ok := ctx.Value(ctxKey).(*TimeBuckets)
    62  	if ok && v != nil {
    63  		return ctx, v
    64  	}
    65  	buckets := NewTimeBuckets(clock, log)
    66  	ctx = context.WithValue(ctx, ctxKey, buckets)
    67  	return ctx, buckets
    68  }