package birate: 支持配置 bitrate 单位，支持由外部传入 unix 时间戳

pkg/bitrate/bitrate.go

@@ -13,12 +13,6 @@ import (
 	"time"
 )
 
-// 先来个最简单的，后续再精细化配置：
-//
-// - 收包时间目前只能由内部获取当前时间，应提供接口支持外部传入
-// - 返回的rate单位固定为 kbit/s
-// - 不需要存储Time结构体，存毫秒级的 unix 时间戳
-
 type Bitrate struct {
 	option Option
 
@@ -26,17 +20,29 @@ type Bitrate struct {
 	bucketSlice []bucket
 }
 
+type Unit uint8
+
+const (
+	UnitBitPerSec Unit = iota + 1
+	UnitBytePerSec
+	UnitKBitPerSec
+	UnitKBytePerSec
+)
+
+// TODO chef: 考虑支持配置是否在内部使用锁
 type Option struct {
 	WindowMS int
+	Unit     Unit
 }
 
 var defaultOption = Option{
 	WindowMS: 1000,
+	Unit:     UnitKBitPerSec,
 }
 
 type bucket struct {
 	n int
-	t time.Time
+	t int64 // unix 时间戳，单位毫秒
 }
 
 type ModOption func(option *Option)
@@ -51,22 +57,33 @@ func NewBitrate(modOptions ...ModOption) *Bitrate {
 	}
 }
 
-func (b *Bitrate) Add(bytes int) {
-	now := time.Now()
+// @param nowUnixMSec 可选择从外部传入当前 unix 时间戳，单位毫秒
+func (b *Bitrate) Add(bytes int, nowUnixMSec ...int64) {
+	var now int64
+	if len(nowUnixMSec) == 0 {
+		now = time.Now().UnixNano() / 1e6
+	} else {
+		now = nowUnixMSec[0]
+	}
+
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
 	b.sweepStale(now)
-
 	b.bucketSlice = append(b.bucketSlice, bucket{
 		n: bytes,
 		t: now,
 	})
 }
 
-// @return 返回值单位 kbit/s
-func (b *Bitrate) Rate() int {
-	now := time.Now()
+func (b *Bitrate) Rate(nowUnixMSec ...int64) float32 {
+	var now int64
+	if len(nowUnixMSec) == 0 {
+		now = time.Now().UnixNano() / 1e6
+	} else {
+		now = nowUnixMSec[0]
+	}
+
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
@@ -76,13 +93,23 @@ func (b *Bitrate) Rate() int {
 		total += b.bucketSlice[i].n
 	}
 
-	// total * 8 / 1000 * 1000 / b.windowMS
-	return total * 8 / b.option.WindowMS
+	var ret float32
+	switch b.option.Unit {
+	case UnitBitPerSec:
+		ret = float32(total*8*1000) / float32(b.option.WindowMS)
+	case UnitBytePerSec:
+		ret = float32(total*1000) / float32(b.option.WindowMS)
+	case UnitKBitPerSec:
+		ret = float32(total*8) / float32(b.option.WindowMS)
+	case UnitKBytePerSec:
+		ret = float32(total) / float32(b.option.WindowMS)
+	}
+	return ret
 }
 
-func (b *Bitrate) sweepStale(now time.Time) {
+func (b *Bitrate) sweepStale(now int64) {
 	for i := range b.bucketSlice {
-		if now.Sub(b.bucketSlice[i].t) > time.Duration(b.option.WindowMS)*time.Millisecond {
+		if now-b.bucketSlice[i].t > int64(b.option.WindowMS) {
 			b.bucketSlice = b.bucketSlice[1:]
 		} else {
 			break

pkg/bitrate/bitrate_test.go

@@ -23,7 +23,36 @@ func TestBitrate(t *testing.T) {
 	})
 	b.Add(1000)
 	r := b.Rate()
-	assert.Equal(t, 800, r)
+	assert.Equal(t, float32(800), r)
 	time.Sleep(100 * time.Millisecond)
 	b.Rate()
 }
+
+func TestUnit(t *testing.T) {
+	golden := map[bitrate.Unit]float32{
+		bitrate.UnitBitPerSec:   800 * 1000,
+		bitrate.UnitBytePerSec:  100 * 1000,
+		bitrate.UnitKBitPerSec:  800,
+		bitrate.UnitKBytePerSec: 100,
+	}
+	for k, v := range golden {
+		b := bitrate.NewBitrate(func(option *bitrate.Option) {
+			option.WindowMS = 10
+			option.Unit = k
+		})
+		b.Add(1000)
+		r := b.Rate()
+		assert.Equal(t, v, r)
+	}
+}
+
+func TestOutsizeNow(t *testing.T) {
+	var b *bitrate.Bitrate
+	b = bitrate.NewBitrate(func(option *bitrate.Option) {
+		option.WindowMS = 10
+	})
+	now := time.Now().UnixNano() / 1e6
+	b.Add(1000, now)
+	r := b.Rate(now)
+	assert.Equal(t, float32(800), r)
+}

pkg/consistenthash/consistenthash.go

@@ -16,6 +16,8 @@ import (
 	"strconv"
 )
 
+var ErrIsEmpty = errors.New("naza.consistenthash: is empty")
+
 type ConsistentHash interface {
 	Add(nodes ...string)
 	Del(nodes ...string)
@@ -29,8 +31,6 @@ type ConsistentHash interface {
 	Nodes() map[string]uint64
 }
 
-var ErrIsEmpty = errors.New("naza.consistenthash: is empty")
-
 type HashFunc func([]byte) uint32
 
 type Option struct {