github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/requester/rio/speeds/speeds.go (about)

     1  // Package speeds 速度计算工具包
     2  package speeds
     3  
     4  import (
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  )
     9  
    10  type (
    11  	// Speeds 统计速度
    12  	Speeds struct {
    13  		count    int64
    14  		interval time.Duration // 刷新周期
    15  		nowTime  time.Time
    16  		once     sync.Once
    17  	}
    18  )
    19  
    20  func (sps *Speeds) initOnce() {
    21  	sps.once.Do(func() {
    22  		sps.nowTime = time.Now()
    23  		if sps.interval <= 0 {
    24  			sps.interval = 1 * time.Second
    25  		}
    26  	})
    27  }
    28  
    29  // SetInterval 设置刷新周期
    30  func (sps *Speeds) SetInterval(interval time.Duration) {
    31  	if interval <= 0 {
    32  		return
    33  	}
    34  	sps.interval = interval
    35  }
    36  
    37  // Add 原子操作, 增加数据量
    38  func (sps *Speeds) Add(count int64) {
    39  	// 初始化
    40  	sps.initOnce()
    41  	atomic.AddInt64(&sps.count, count)
    42  }
    43  
    44  // GetSpeeds 结束统计速度, 并返回速度
    45  func (sps *Speeds) GetSpeeds() (speeds int64) {
    46  	sps.initOnce()
    47  
    48  	since := time.Since(sps.nowTime)
    49  	if since <= 0 {
    50  		return 0
    51  	}
    52  	speeds = int64(float64(atomic.LoadInt64(&sps.count)) * sps.interval.Seconds() / since.Seconds())
    53  
    54  	// 更新下一轮
    55  	if since >= sps.interval {
    56  		atomic.StoreInt64(&sps.count, 0)
    57  		sps.nowTime = time.Now()
    58  	}
    59  	return
    60  }