Initial commit

library/stat/counter/BUILD

@@ -0,0 +1,45 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = [
+        "counter_test.go",
+        "rolling_test.go",
+    ],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = ["//vendor/github.com/stretchr/testify/assert:go_default_library"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "counter.go",
+        "gauge.go",
+        "rolling.go",
+    ],
+    importpath = "go-common/library/stat/counter",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/stat/counter/counter.go

@@ -0,0 +1,62 @@
+package counter
+
+import (
+	"sync"
+)
+
+// Counter is a counter interface.
+type Counter interface {
+	Add(int64)
+	Reset()
+	Value() int64
+}
+
+// Group is a counter group.
+type Group struct {
+	mu   sync.RWMutex
+	vecs map[string]Counter
+
+	// New optionally specifies a function to generate a counter.
+	// It may not be changed concurrently with calls to other functions.
+	New func() Counter
+}
+
+// Add add a counter by a specified key, if counter not exists then make a new one and return new value.
+func (g *Group) Add(key string, value int64) {
+	g.mu.RLock()
+	vec, ok := g.vecs[key]
+	g.mu.RUnlock()
+	if !ok {
+		vec = g.New()
+		g.mu.Lock()
+		if g.vecs == nil {
+			g.vecs = make(map[string]Counter)
+		}
+		if _, ok = g.vecs[key]; !ok {
+			g.vecs[key] = vec
+		}
+		g.mu.Unlock()
+	}
+	vec.Add(value)
+}
+
+// Value get a counter value by key.
+func (g *Group) Value(key string) int64 {
+	g.mu.RLock()
+	vec, ok := g.vecs[key]
+	g.mu.RUnlock()
+	if ok {
+		return vec.Value()
+	}
+	return 0
+}
+
+// Reset reset a counter by key.
+func (g *Group) Reset(key string) {
+	g.mu.RLock()
+	vec, ok := g.vecs[key]
+	g.mu.RUnlock()
+	if ok {
+		vec.Reset()
+	}
+}

library/stat/counter/counter_test.go

@@ -0,0 +1,57 @@
+package counter
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGaugeCounter(t *testing.T) {
+	key := "test"
+	g := &Group{
+		New: func() Counter {
+			return NewGauge()
+		},
+	}
+	g.Add(key, 1)
+	g.Add(key, 2)
+	g.Add(key, 3)
+	g.Add(key, -1)
+	assert.Equal(t, g.Value(key), int64(5))
+	g.Reset(key)
+	assert.Equal(t, g.Value(key), int64(0))
+}
+
+func TestRollingCounter(t *testing.T) {
+	key := "test"
+	g := &Group{
+		New: func() Counter {
+			return NewRolling(time.Second, 10)
+		},
+	}
+
+	t.Run("add_key_b1", func(t *testing.T) {
+		g.Add(key, 1)
+		assert.Equal(t, g.Value(key), int64(1))
+	})
+	time.Sleep(time.Millisecond * 110)
+	t.Run("add_key_b2", func(t *testing.T) {
+		g.Add(key, 1)
+		assert.Equal(t, g.Value(key), int64(2))
+	})
+	time.Sleep(time.Millisecond * 900) // expire one bucket, 110 + 900
+	t.Run("expire_b1", func(t *testing.T) {
+		assert.Equal(t, g.Value(key), int64(1))
+		g.Add(key, 1)
+		assert.Equal(t, g.Value(key), int64(2)) // expire one bucket
+	})
+	time.Sleep(time.Millisecond * 1100)
+	t.Run("expire_all", func(t *testing.T) {
+		assert.Equal(t, g.Value(key), int64(0))
+	})
+	t.Run("reset", func(t *testing.T) {
+		g.Reset(key)
+		assert.Equal(t, g.Value(key), int64(0))
+	})
+}

library/stat/counter/gauge.go

@@ -0,0 +1,28 @@
+package counter
+
+import "sync/atomic"
+
+var _ Counter = new(gaugeCounter)
+
+// A value is a thread-safe counter implementation.
+type gaugeCounter int64
+
+// NewGauge return a guage counter.
+func NewGauge() Counter {
+	return new(gaugeCounter)
+}
+
+// Add method increments the counter by some value and return new value
+func (v *gaugeCounter) Add(val int64) {
+	atomic.AddInt64((*int64)(v), val)
+}
+
+// Value method returns the counter's current value.
+func (v *gaugeCounter) Value() int64 {
+	return atomic.LoadInt64((*int64)(v))
+}
+
+// Reset reset the counter.
+func (v *gaugeCounter) Reset() {
+	atomic.StoreInt64((*int64)(v), 0)
+}

library/stat/counter/rolling.go

@@ -0,0 +1,117 @@
+package counter
+
+import (
+	"sync"
+	"time"
+)
+
+type bucket struct {
+	val  int64
+	next *bucket
+}
+
+func (b *bucket) Add(val int64) {
+	b.val += val
+}
+
+func (b *bucket) Value() int64 {
+	return b.val
+}
+
+func (b *bucket) Reset() {
+	b.val = 0
+}
+
+var _ Counter = new(rollingCounter)
+
+type rollingCounter struct {
+	mu         sync.RWMutex
+	buckets    []bucket
+	bucketTime int64
+	lastAccess int64
+	cur        *bucket
+}
+
+// NewRolling creates a new window. windowTime is the time covering the entire
+// window. windowBuckets is the number of buckets the window is divided into.
+// An example: a 10 second window with 10 buckets will have 10 buckets covering
+// 1 second each.
+func NewRolling(window time.Duration, winBucket int) Counter {
+	buckets := make([]bucket, winBucket)
+	bucket := &buckets[0]
+	for i := 1; i < winBucket; i++ {
+		bucket.next = &buckets[i]
+		bucket = bucket.next
+	}
+	bucket.next = &buckets[0]
+	bucketTime := time.Duration(window.Nanoseconds() / int64(winBucket))
+	return &rollingCounter{
+		cur:        &buckets[0],
+		buckets:    buckets,
+		bucketTime: int64(bucketTime),
+		lastAccess: time.Now().UnixNano(),
+	}
+}
+
+// Add increments the counter by value and return new value.
+func (r *rollingCounter) Add(val int64) {
+	r.mu.Lock()
+	r.lastBucket().Add(val)
+	r.mu.Unlock()
+}
+
+// Value get the counter value.
+func (r *rollingCounter) Value() (sum int64) {
+	now := time.Now().UnixNano()
+	r.mu.RLock()
+	b := r.cur
+	i := r.elapsed(now)
+	for j := 0; j < len(r.buckets); j++ {
+		// skip all future reset bucket.
+		if i > 0 {
+			i--
+		} else {
+			sum += b.Value()
+		}
+		b = b.next
+	}
+	r.mu.RUnlock()
+	return
+}
+
+//  Reset reset the counter.
+func (r *rollingCounter) Reset() {
+	r.mu.Lock()
+	for i := range r.buckets {
+		r.buckets[i].Reset()
+	}
+	r.mu.Unlock()
+}
+
+func (r *rollingCounter) elapsed(now int64) (i int) {
+	var e int64
+	if e = now - r.lastAccess; e <= r.bucketTime {
+		return
+	}
+	if i = int(e / r.bucketTime); i > len(r.buckets) {
+		i = len(r.buckets)
+	}
+	return
+}
+
+func (r *rollingCounter) lastBucket() (b *bucket) {
+	now := time.Now().UnixNano()
+	b = r.cur
+	// reset the buckets between now and number of buckets ago. If
+	// that is more that the existing buckets, reset all.
+	if i := r.elapsed(now); i > 0 {
+		r.lastAccess = now
+		for ; i > 0; i-- {
+			// replace the next used bucket.
+			b = b.next
+			b.Reset()
+		}
+	}
+	r.cur = b
+	return
+}

library/stat/counter/rolling_test.go

@@ -0,0 +1,23 @@
+package counter
+
+import (
+	"testing"
+	"time"
+)
+
+func TestRollingCounterMinInterval(t *testing.T) {
+	count := NewRolling(time.Second/2, 10)
+	tk1 := time.NewTicker(5 * time.Millisecond)
+	defer tk1.Stop()
+	for i := 0; i < 100; i++ {
+		<-tk1.C
+		count.Add(1)
+	}
+
+	v := count.Value()
+	t.Logf("count value: %d", v)
+	// 10% of error when bucket is 10
+	if v < 90 || v > 110 {
+		t.Errorf("expect value in [90-110] get %d", v)
+	}
+}