1package profiling
2
3import (
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
12type ctxKeyType string
13
14var ctxKey = ctxKeyType("timebuckets")
15
16type TimeBuckets struct {
17	sync.Mutex
18	clock clockwork.Clock
19	log   logger.Logger
20	times map[string]time.Duration
21}
22
23func 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
31func (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
41func (t *TimeBuckets) Get(bucketName string) time.Duration {
42	t.Lock()
43	defer t.Unlock()
44	return t.times[bucketName]
45}
46
47func (t *TimeBuckets) Log(ctx context.Context, bucketName string) {
48	t.log.CDebugf(ctx, "TimeBucket %s [time=%s]", bucketName, t.Get(bucketName))
49}
50
51func (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
58type FinFn func()
59
60func 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}
69