JCS-pub/client/internal/downloader/strategy/selector.go

package strategy

import (
	"fmt"
	"math"
	"reflect"

	"github.com/samber/lo"
	"gitlink.org.cn/cloudream/common/pkgs/bitmap"
	"gitlink.org.cn/cloudream/common/utils/math2"
	"gitlink.org.cn/cloudream/common/utils/sort2"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/metacache"
	"gitlink.org.cn/cloudream/jcs-pub/common/consts"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

type Request struct {
	Detail       jcstypes.ObjectDetail
	Range        math2.Range
	DestLocation jcstypes.Location
}

type Strategy interface {
	GetDetail() jcstypes.ObjectDetail
}

// 直接下载完整对象
type DirectStrategy struct {
	Detail    jcstypes.ObjectDetail
	UserSpace jcstypes.UserSpaceDetail
}

func (s *DirectStrategy) GetDetail() jcstypes.ObjectDetail {
	return s.Detail
}

// 从指定对象重建对象
type ECReconstructStrategy struct {
	Detail     jcstypes.ObjectDetail
	Redundancy jcstypes.ECRedundancy
	Blocks     []jcstypes.ObjectBlock
	UserSpaces []jcstypes.UserSpaceDetail
}

func (s *ECReconstructStrategy) GetDetail() jcstypes.ObjectDetail {
	return s.Detail
}

type LRCReconstructStrategy struct {
	Detail     jcstypes.ObjectDetail
	Redundancy jcstypes.LRCRedundancy
	Blocks     []jcstypes.ObjectBlock
	Spaces     []jcstypes.UserSpaceDetail
}

func (s *LRCReconstructStrategy) GetDetail() jcstypes.ObjectDetail {
	return s.Detail
}

type Selector struct {
	cfg          Config
	spaceMeta    *metacache.UserSpaceMeta
	hubMeta      *metacache.HubMeta
	connectivity *metacache.Connectivity
}

func NewSelector(cfg Config, storageMeta *metacache.UserSpaceMeta, hubMeta *metacache.HubMeta, connectivity *metacache.Connectivity) *Selector {
	return &Selector{
		cfg:          cfg,
		spaceMeta:    storageMeta,
		hubMeta:      hubMeta,
		connectivity: connectivity,
	}
}

func (s *Selector) Select(req Request) (Strategy, error) {
	req2 := request2{
		Detail:       req.Detail,
		Range:        req.Range,
		DestLocation: req.DestLocation,
	}

	switch red := req.Detail.Object.Redundancy.(type) {
	case *jcstypes.NoneRedundancy:
		return s.selectForNoneOrRep(req2)

	case *jcstypes.RepRedundancy:
		return s.selectForNoneOrRep(req2)

	case *jcstypes.ECRedundancy:
		return s.selectForEC(req2, *red)

	case *jcstypes.LRCRedundancy:
		return s.selectForLRC(req2, *red)
	}

	return nil, fmt.Errorf("unsupported redundancy type: %v of object %v", reflect.TypeOf(req.Detail.Object.Redundancy), req.Detail.Object.ObjectID)
}

type downloadSpaceInfo struct {
	Space        jcstypes.UserSpaceDetail
	ObjectPinned bool
	Blocks       []jcstypes.ObjectBlock
	Distance     float64
}

type downloadBlock struct {
	Space jcstypes.UserSpaceDetail
	Block jcstypes.ObjectBlock
}

type request2 struct {
	Detail       jcstypes.ObjectDetail
	Range        math2.Range
	DestLocation jcstypes.Location
}

func (s *Selector) selectForNoneOrRep(req request2) (Strategy, error) {
	sortedStgs := s.sortDownloadStorages(req)
	if len(sortedStgs) == 0 {
		return nil, fmt.Errorf("no storage available for download")
	}

	_, blks := s.getMinReadingBlockSolution(sortedStgs, 1)
	if len(blks) == 0 {
		return nil, fmt.Errorf("no block available for download")
	}

	return &DirectStrategy{
		Detail:    req.Detail,
		UserSpace: sortedStgs[0].Space,
	}, nil
}

func (s *Selector) selectForEC(req request2, red jcstypes.ECRedundancy) (Strategy, error) {
	sortedStgs := s.sortDownloadStorages(req)
	if len(sortedStgs) == 0 {
		return nil, fmt.Errorf("no storage available for download")
	}

	bsc, blocks := s.getMinReadingBlockSolution(sortedStgs, red.K)
	osc, stg := s.getMinReadingObjectSolution(sortedStgs, red.K)

	if bsc < osc {
		bs := make([]jcstypes.ObjectBlock, len(blocks))
		ss := make([]jcstypes.UserSpaceDetail, len(blocks))
		for i, b := range blocks {
			bs[i] = b.Block
			ss[i] = b.Space
		}

		return &ECReconstructStrategy{
			Detail:     req.Detail,
			Redundancy: red,
			Blocks:     bs,
			UserSpaces: ss,
		}, nil
	}

	// bsc >= osc，如果osc是MaxFloat64，那么bsc也一定是，也就意味着没有足够块来恢复文件
	if osc == math.MaxFloat64 {
		return nil, fmt.Errorf("no enough blocks to reconstruct the object %v , want %d, get only %d", req.Detail.Object.ObjectID, red.K, len(blocks))
	}

	return &DirectStrategy{
		Detail:    req.Detail,
		UserSpace: stg,
	}, nil
}

func (s *Selector) selectForLRC(req request2, red jcstypes.LRCRedundancy) (Strategy, error) {
	sortedStgs := s.sortDownloadStorages(req)
	if len(sortedStgs) == 0 {
		return nil, fmt.Errorf("no storage available for download")
	}

	var blocks []downloadBlock
	selectedBlkIdx := make(map[int]bool)
	for _, stg := range sortedStgs {
		for _, b := range stg.Blocks {
			if b.Index >= red.M() || selectedBlkIdx[b.Index] {
				continue
			}
			blocks = append(blocks, downloadBlock{
				Space: stg.Space,
				Block: b,
			})
			selectedBlkIdx[b.Index] = true
		}
	}
	if len(blocks) < red.K {
		return nil, fmt.Errorf("not enough blocks to download lrc object")
	}

	bs := make([]jcstypes.ObjectBlock, len(blocks))
	ss := make([]jcstypes.UserSpaceDetail, len(blocks))
	for i, b := range blocks {
		bs[i] = b.Block
		ss[i] = b.Space
	}

	return &LRCReconstructStrategy{
		Detail:     req.Detail,
		Redundancy: red,
		Blocks:     bs,
		Spaces:     ss,
	}, nil
}

func (s *Selector) sortDownloadStorages(req request2) []*downloadSpaceInfo {
	var spaceIDs []jcstypes.UserSpaceID
	for _, id := range req.Detail.PinnedAt {
		if !lo.Contains(spaceIDs, id) {
			spaceIDs = append(spaceIDs, id)
		}
	}
	for _, b := range req.Detail.Blocks {
		if !lo.Contains(spaceIDs, b.UserSpaceID) {
			spaceIDs = append(spaceIDs, b.UserSpaceID)
		}
	}

	downloadSpaceMap := make(map[jcstypes.UserSpaceID]*downloadSpaceInfo)
	for _, id := range req.Detail.PinnedAt {
		storage, ok := downloadSpaceMap[id]
		if !ok {
			mod := s.spaceMeta.Get(id)
			if mod == nil {
				continue
			}

			storage = &downloadSpaceInfo{
				Space:        *mod,
				ObjectPinned: true,
				Distance:     s.getStorageDistance(req, *mod),
			}
			downloadSpaceMap[id] = storage
		}

		storage.ObjectPinned = true
	}

	for _, b := range req.Detail.Blocks {
		space, ok := downloadSpaceMap[b.UserSpaceID]
		if !ok {
			mod := s.spaceMeta.Get(b.UserSpaceID)
			if mod == nil {
				continue
			}

			space = &downloadSpaceInfo{
				Space:    *mod,
				Distance: s.getStorageDistance(req, *mod),
			}
			downloadSpaceMap[b.UserSpaceID] = space
		}

		space.Blocks = append(space.Blocks, b)
	}

	return sort2.Sort(lo.Values(downloadSpaceMap), func(left, right *downloadSpaceInfo) int {
		return sort2.Cmp(left.Distance, right.Distance)
	})
}

func (s *Selector) getStorageDistance(req request2, src jcstypes.UserSpaceDetail) float64 {
	// TODO 重新设计计算方式

	// if req.DestHub != nil {
	// 	if src.RecommendHub.HubID == req.DestHub.HubID {
	// 		return consts.StorageDistanceSameStorage
	// 	}

	// 	if src.RecommendHub.LocationID == req.DestHub.LocationID {
	// 		return consts.StorageDistanceSameLocation
	// 	}

	// 	latency := s.connectivity.Get(src.RecommendHub.HubID, req.DestHub.HubID)
	// 	if latency == nil || *latency > time.Duration(float64(time.Millisecond)*s.cfg.HighLatencyHubMs) {
	// 		return consts.HubDistanceHighLatencyHub
	// 	}

	// 	return consts.StorageDistanceOther
	// }

	if src.UserSpace.Storage.GetLocation() == req.DestLocation {
		return consts.StorageDistanceSameStorage
	}

	return consts.StorageDistanceOther
}

func (s *Selector) getMinReadingBlockSolution(sortedStgs []*downloadSpaceInfo, k int) (float64, []downloadBlock) {
	gotBlocksMap := bitmap.Bitmap64(0)
	var gotBlocks []downloadBlock
	dist := float64(0.0)
	for _, n := range sortedStgs {
		for _, b := range n.Blocks {
			if !gotBlocksMap.Get(b.Index) {
				gotBlocks = append(gotBlocks, downloadBlock{
					Space: n.Space,
					Block: b,
				})
				gotBlocksMap.Set(b.Index, true)
				dist += n.Distance
			}

			if len(gotBlocks) >= k {
				return dist, gotBlocks
			}
		}
	}

	return math.MaxFloat64, gotBlocks
}

func (s *Selector) getMinReadingObjectSolution(sortedStgs []*downloadSpaceInfo, k int) (float64, jcstypes.UserSpaceDetail) {
	dist := math.MaxFloat64
	var downloadSpace jcstypes.UserSpaceDetail
	for _, n := range sortedStgs {
		if n.ObjectPinned && float64(k)*n.Distance < dist {
			dist = float64(k) * n.Distance
			stg := n.Space
			downloadSpace = stg
		}
	}

	return dist, downloadSpace
}