JCS-pub/client/internal/speedstats/speedstats.go

package speedstats

import (
	"fmt"
	"math/rand/v2"
	"sync"
	"time"
)

const (
	LastStatsWeigth = 0.8
	MinProb         = 0.05
)

type Stats struct {
	TotalSize    int64
	TotalTime    time.Duration
	LastSpeed    float64
	HasLastSpeed bool
	AvarageSpeed float64
}

func (s *Stats) Record(size int64, time time.Duration) {
	s.TotalSize += size
	s.TotalTime += time

	if s.HasLastSpeed {
		s.AvarageSpeed = float64(s.TotalSize)*(1-LastStatsWeigth)/s.TotalTime.Seconds() + s.LastSpeed*LastStatsWeigth
	} else {
		s.AvarageSpeed = float64(s.TotalSize) / s.TotalTime.Seconds()
	}
}

func (s *Stats) Step() {
	s.TotalSize = 0
	s.TotalTime = 0
	s.LastSpeed = s.AvarageSpeed
	s.AvarageSpeed = s.LastSpeed * LastStatsWeigth
	s.HasLastSpeed = true
}

type SpeedStats struct {
	lock         sync.RWMutex
	stats100M    []Stats
	stats1G      []Stats
	statsAbove1G []Stats
}

func New() *SpeedStats {
	return &SpeedStats{
		stats100M:    make([]Stats, 2),
		stats1G:      make([]Stats, 2),
		statsAbove1G: make([]Stats, 2),
	}
}

func (p *SpeedStats) Record(size int64, time time.Duration, happenedAtClient bool) {
	p.lock.Lock()
	defer p.lock.Unlock()

	if size < 100*1024*1024 {
		if happenedAtClient {
			p.stats100M[0].Record(size, time)
		} else {
			p.stats100M[1].Record(size, time)
		}
	} else if size < 1024*1024*1024 {
		if happenedAtClient {
			p.stats1G[0].Record(size, time)
		} else {
			p.stats1G[1].Record(size, time)
		}
	} else {
		if happenedAtClient {
			p.statsAbove1G[0].Record(size, time)
		} else {
			p.statsAbove1G[1].Record(size, time)
		}
	}
}

func (p *SpeedStats) Step() {
	p.lock.Lock()
	defer p.lock.Unlock()

	p.stats100M[0].Step()
	p.stats100M[1].Step()
	p.stats1G[0].Step()
	p.stats1G[1].Step()
	p.statsAbove1G[0].Step()
	p.statsAbove1G[1].Step()
}

func (p *SpeedStats) Dump() string {
	p.lock.RLock()
	defer p.lock.RUnlock()

	return fmt.Sprintf("100M: %v, %v\n1G: %v, %v\nAbove1G: %v, %v\n", p.stats100M[0].AvarageSpeed, p.stats100M[1].AvarageSpeed, p.stats1G[0].AvarageSpeed, p.stats1G[1].AvarageSpeed, p.statsAbove1G[0].AvarageSpeed, p.statsAbove1G[1].AvarageSpeed)
}

func (p *SpeedStats) ShouldAtClient(size int64) bool {
	p.lock.RLock()
	defer p.lock.RUnlock()

	var ss []Stats

	if size < 100*1024*1024 {
		ss = p.stats100M
	} else if size < 1024*1024*1024 {
		ss = p.stats1G
	} else {
		ss = p.statsAbove1G
	}

	prob := 0.0

	if ss[0].AvarageSpeed == 0 || ss[1].AvarageSpeed == 0 {
		prob = 0.5
	} else {
		// 保证最小概率为0.05，最大概率为0.95
		totalSpeed := ss[0].AvarageSpeed + ss[1].AvarageSpeed
		scale := 1 - 2*MinProb
		prob = ss[0].AvarageSpeed/totalSpeed*scale + MinProb
	}

	v := rand.Float64()
	return v < prob
}