JCS-pub/common/pkgs/servicestats/hub_strorage_transfer.go

package servicestats

import (
	"math"
	"sync"
	"time"

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

type HubStorageTransferStats struct {
	data      HubStorageTransferStatsData
	fromHubID jcstypes.HubID
	lock      *sync.Mutex
}

type HubStorageTransferStatsData struct {
	Entries   map[jcstypes.StorageID]*HubStorageTransferStatsEntry
	StartTime time.Time
	EndTime   time.Time
}

type HubStorageTransferStatsEntry struct {
	DestStorageID jcstypes.StorageID

	OutputBytes       int64
	MaxOutputBytes    int64
	MinOutputBytes    int64
	TotalOutputReqs   int64
	SuccessOutputReqs int64

	InputBytes       int64
	MaxInputBytes    int64
	MinInputBytes    int64
	TotalInputReqs   int64
	SuccessInputReqs int64
}

func (s *HubStorageTransferStats) RecordUpload(dstStorageID jcstypes.StorageID, transferBytes int64, isSuccess bool) {
	s.lock.Lock()
	defer s.lock.Unlock()

	e := s.data.Entries[dstStorageID]
	if e == nil {
		e = &HubStorageTransferStatsEntry{
			DestStorageID:  dstStorageID,
			MinInputBytes:  math.MaxInt64,
			MinOutputBytes: math.MaxInt64,
		}
		s.data.Entries[dstStorageID] = e
	}
	e.OutputBytes += transferBytes
	e.MaxOutputBytes = math2.Max(e.MaxOutputBytes, transferBytes)
	e.MinOutputBytes = math2.Min(e.MinOutputBytes, transferBytes)
	if isSuccess {
		e.SuccessOutputReqs++
	}
	e.TotalOutputReqs++
}

func (s *HubStorageTransferStats) RecordDownload(dstStorageID jcstypes.StorageID, transferBytes int64, isSuccess bool) {
	s.lock.Lock()
	defer s.lock.Unlock()

	e := s.data.Entries[dstStorageID]
	if e == nil {
		e = &HubStorageTransferStatsEntry{
			DestStorageID:  dstStorageID,
			MinInputBytes:  math.MaxInt64,
			MinOutputBytes: math.MaxInt64,
		}
		s.data.Entries[dstStorageID] = e
	}
	e.InputBytes += transferBytes
	e.MaxInputBytes = math2.Max(e.MaxInputBytes, transferBytes)
	e.MinInputBytes = math2.Min(e.MinInputBytes, transferBytes)
	if isSuccess {
		e.SuccessInputReqs++
	}
}

func (s *HubStorageTransferStats) DumpData(reset bool) HubStorageTransferStatsData {
	s.lock.Lock()
	defer s.lock.Unlock()

	data := s.data
	data.Entries = make(map[jcstypes.StorageID]*HubStorageTransferStatsEntry)
	for k, v := range s.data.Entries {
		v2 := *v
		data.Entries[k] = &v2
	}
	data.EndTime = time.Now()

	if reset {
		s.data.Entries = make(map[jcstypes.StorageID]*HubStorageTransferStatsEntry)
		s.data.StartTime = time.Now()
	}

	return data
}