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

package speedstats

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

	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

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, elapsedTime time.Duration, happenedAtClient bool) {
	p.lock.Lock()
	defer p.lock.Unlock()

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

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) 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
}

func (p *SpeedStats) DumpStatus() jcstypes.SpeedStatsStatus {
	return jcstypes.SpeedStatsStatus{
		Below100M: []jcstypes.SpeedStatsStatusEntry{
			{
				TotalSize:    p.stats100M[0].TotalSize,
				TotalTime:    p.stats100M[0].TotalTime,
				AvarageSpeed: p.stats100M[0].AvarageSpeed,
				LastSpeed:    p.stats100M[0].LastSpeed,
			},
			{
				TotalSize:    p.stats100M[1].TotalSize,
				TotalTime:    p.stats100M[1].TotalTime,
				AvarageSpeed: p.stats100M[1].AvarageSpeed,
				LastSpeed:    p.stats100M[1].LastSpeed,
			},
		},
		Below1G: []jcstypes.SpeedStatsStatusEntry{
			{
				TotalSize:    p.stats1G[0].TotalSize,
				TotalTime:    p.stats1G[0].TotalTime,
				AvarageSpeed: p.stats1G[0].AvarageSpeed,
				LastSpeed:    p.stats1G[0].LastSpeed,
			},
			{
				TotalSize:    p.stats1G[1].TotalSize,
				TotalTime:    p.stats1G[1].TotalTime,
				AvarageSpeed: p.stats1G[1].AvarageSpeed,
				LastSpeed:    p.stats1G[1].LastSpeed,
			},
		},
		Above1G: []jcstypes.SpeedStatsStatusEntry{
			{
				TotalSize:    p.statsAbove1G[0].TotalSize,
				TotalTime:    p.statsAbove1G[0].TotalTime,
				AvarageSpeed: p.statsAbove1G[0].AvarageSpeed,
				LastSpeed:    p.statsAbove1G[0].LastSpeed,
			},
			{
				TotalSize:    p.statsAbove1G[1].TotalSize,
				TotalTime:    p.statsAbove1G[1].TotalTime,
				AvarageSpeed: p.statsAbove1G[1].AvarageSpeed,
				LastSpeed:    p.statsAbove1G[1].LastSpeed,
			},
		},
	}
}