naza/pkg/bitrate/bitrate.go

// Copyright 2019, Chef.  All rights reserved.
// https://github.com/q191201771/naza
//
// Use of this source code is governed by a MIT-style license
// that can be found in the License file.
//
// Author: Chef (191201771@qq.com)

package bitrate

import (
	"sync"
	"time"
)

// 先来个最简单的，后续再精细化配置：
//
// - 收包时间目前只能由内部获取当前时间，应提供接口支持外部传入
// - 返回的rate单位固定为 kbit/s
// - 不需要存储Time结构体，存毫秒级的 unix 时间戳

type Bitrate struct {
	option Option

	mu          sync.Mutex
	bucketSlice []bucket
}

type Option struct {
	WindowMS int
}

var defaultOption = Option{
	WindowMS: 1000,
}

type bucket struct {
	n int
	t time.Time
}

type ModOption func(option *Option)

func NewBitrate(modOptions ...ModOption) *Bitrate {
	option := defaultOption
	for _, fn := range modOptions {
		fn(&option)
	}
	return &Bitrate{
		option: option,
	}
}

func (b *Bitrate) Add(bytes int) {
	now := time.Now()
	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()
	b.mu.Lock()
	defer b.mu.Unlock()

	b.sweepStale(now)
	var total int
	for i := range b.bucketSlice {
		total += b.bucketSlice[i].n
	}

	// total * 8 / 1000 * 1000 / b.windowMS
	return total * 8 / b.option.WindowMS
}

func (b *Bitrate) sweepStale(now time.Time) {
	for i := range b.bucketSlice {
		if now.Sub(b.bucketSlice[i].t) > time.Duration(b.option.WindowMS)*time.Millisecond {
			b.bucketSlice = b.bucketSlice[1:]
		} else {
			break
		}
	}
}