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

package servicestats

import (
	"math"
	"sync"
	"time"

	"gitlink.org.cn/cloudream/common/utils/math2"
	cortypes "gitlink.org.cn/cloudream/storage2/coordinator/types"
)

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

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

type HubStorageTransferStatsEntry struct {
	DestStorageID cortypes.StorageID

	OutputBytes    int64
	MaxOutputBytes int64
	MinOutputBytes int64
	TotalOutput    int64
	SuccessOutput  int64

	InputBytes    int64
	MaxInputBytes int64
	MinInputBytes int64
	TotalInput    int64
	SuccessInput  int64
}

func (s *HubStorageTransferStats) RecordUpload(dstStorageID cortypes.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.SuccessOutput++
	}
	e.TotalOutput++
}

func (s *HubStorageTransferStats) RecordDownload(dstStorageID cortypes.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.SuccessInput++
	}
}

func (s *HubStorageTransferStats) Reset() time.Time {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.data.Entries = make(map[cortypes.StorageID]*HubStorageTransferStatsEntry)
	s.data.StartTime = time.Now()
	return s.data.StartTime
}

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

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

	return data
}