JCS-pub/client/internal/downloader/iterator.go

package downloader

import (
	"context"
	"fmt"
	"io"
	"reflect"

	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"

	"gitlink.org.cn/cloudream/common/utils/io2"
	"gitlink.org.cn/cloudream/common/utils/math2"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/downloader/strategy"
	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/ops2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2/parser"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/iterator"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/publock"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

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

type DownloadContext struct {
	PubLock *publock.Service
}
type DownloadObjectIterator struct {
	OnClosing    func()
	downloader   *Downloader
	reqs         []downloadReqeust2
	currentIndex int
}

func NewDownloadObjectIterator(downloader *Downloader, downloadObjs []downloadReqeust2) *DownloadObjectIterator {
	return &DownloadObjectIterator{
		downloader: downloader,
		reqs:       downloadObjs,
	}
}

func (i *DownloadObjectIterator) MoveNext() (*Downloading, error) {
	if i.currentIndex >= len(i.reqs) {
		return nil, iterator.ErrNoMoreItem
	}

	req := i.reqs[i.currentIndex]
	if req.Detail == nil {
		return &Downloading{
			Object:  nil,
			File:    nil,
			Request: req.Raw,
		}, nil
	}

	strg, err := i.downloader.selector.Select(strategy.Request{
		Detail:       *req.Detail,
		Range:        math2.NewRange(req.Raw.Offset, req.Raw.Length),
		DestLocation: stgglb.Local.Location,
	})
	if err != nil {
		return nil, fmt.Errorf("selecting download strategy: %w", err)
	}

	var reader io.ReadCloser
	switch strg := strg.(type) {
	case *strategy.DirectStrategy:
		reader, err = i.downloadDirect(req, *strg)
		if err != nil {
			return nil, fmt.Errorf("downloading object %v: %w", req.Raw.ObjectID, err)
		}

	case *strategy.ECReconstructStrategy:
		reader, err = i.downloadECReconstruct(req, *strg)
		if err != nil {
			return nil, fmt.Errorf("downloading ec object %v: %w", req.Raw.ObjectID, err)
		}

	case *strategy.LRCReconstructStrategy:
		reader, err = i.downloadLRCReconstruct(req, *strg)
		if err != nil {
			return nil, fmt.Errorf("downloading lrc object %v: %w", req.Raw.ObjectID, err)
		}

	default:
		return nil, fmt.Errorf("unsupported strategy type: %v", reflect.TypeOf(strg))
	}

	i.currentIndex++
	return &Downloading{
		Object:  &req.Detail.Object,
		File:    reader,
		Request: req.Raw,
	}, nil
}

func (i *DownloadObjectIterator) Close() {
	if i.OnClosing != nil {
		i.OnClosing()
	}
}

func (i *DownloadObjectIterator) downloadDirect(req downloadReqeust2, strg strategy.DirectStrategy) (io.ReadCloser, error) {
	logger.Debugf("downloading object %v from storage %v", req.Raw.ObjectID, strg.UserSpace.UserSpace.Storage.String())

	var strHandle *exec.DriverReadStream
	ft := ioswitch2.NewFromTo()

	toExec, handle := ioswitch2.NewToDriver(ioswitch2.RawStream())
	toExec.Range = math2.Range{
		Offset: req.Raw.Offset,
	}
	len := req.Detail.Object.Size - req.Raw.Offset
	if req.Raw.Length != -1 {
		len = req.Raw.Length
		toExec.Range.Length = &len
	}

	fromSpace := strg.UserSpace

	shouldAtClient := i.downloader.speedStats.ShouldAtClient(len)
	if shouldAtClient {
		fromSpace.RecommendHub = nil
	}

	ft.AddFrom(ioswitch2.NewFromShardstore(req.Detail.Object.FileHash, fromSpace, ioswitch2.RawStream())).AddTo(toExec)
	strHandle = handle

	plans := exec.NewPlanBuilder()
	if err := parser.Parse(ft, plans); err != nil {
		return nil, fmt.Errorf("parsing plan: %w", err)
	}

	exeCtx := exec.NewExecContext()
	exec.SetValueByType(exeCtx, i.downloader.stgPool)
	exec := plans.Execute(exeCtx)
	go func() {
		ret, err := exec.Wait(context.TODO())
		if err != nil {
			logger.Warnf("downloading object %v: %v", req.Raw.ObjectID, err)
		}

		for _, v := range ret.GetArray(ops2.BaseReadStatsStoreKey) {
			v2 := v.(*ops2.BaseReadStatsValue)
			i.downloader.speedStats.Record(v2.Size, v2.Time, v2.Location.IsDriver)
		}
	}()

	return exec.BeginRead(strHandle)
}

func (i *DownloadObjectIterator) downloadECReconstruct(req downloadReqeust2, strg strategy.ECReconstructStrategy) (io.ReadCloser, error) {
	var logStrs []any = []any{fmt.Sprintf("downloading ec object %v from: ", req.Raw.ObjectID)}
	for i, b := range strg.Blocks {
		if i > 0 {
			logStrs = append(logStrs, ", ")
		}

		logStrs = append(logStrs, fmt.Sprintf("%v@%v", b.Index, strg.UserSpaces[i].UserSpace.Storage.String()))
	}
	logger.Debug(logStrs...)

	length := req.Detail.Object.Size - req.Raw.Offset
	if req.Raw.Length != -1 {
		length = req.Raw.Length
	}

	shouldAtClient := i.downloader.speedStats.ShouldAtClient(length)

	downloadBlks := make([]downloadBlock, len(strg.Blocks))
	for i, b := range strg.Blocks {
		fromSpace := strg.UserSpaces[i]
		if shouldAtClient {
			fromSpace.RecommendHub = nil
		}

		downloadBlks[i] = downloadBlock{
			Block: b,
			Space: fromSpace,
		}
	}

	pr, pw := io.Pipe()
	go func() {
		readPos := req.Raw.Offset
		totalReadLen := req.Detail.Object.Size - req.Raw.Offset
		if req.Raw.Length >= 0 {
			totalReadLen = math2.Min(req.Raw.Length, totalReadLen)
		}

		firstStripIndex := readPos / strg.Redundancy.StripSize()
		stripIter := NewStripIterator(i.downloader, req.Detail.Object, downloadBlks, strg.Redundancy, firstStripIndex, i.downloader.strips, i.downloader.cfg.ECStripPrefetchCount)
		defer stripIter.Close()

		for totalReadLen > 0 {
			strip, err := stripIter.MoveNext()
			if err == iterator.ErrNoMoreItem {
				pw.CloseWithError(io.ErrUnexpectedEOF)
				return
			}
			if err != nil {
				pw.CloseWithError(err)
				return
			}

			readRelativePos := readPos - strip.Position
			curReadLen := math2.Min(totalReadLen, strg.Redundancy.StripSize()-readRelativePos)

			err = io2.WriteAll(pw, strip.Data[readRelativePos:readRelativePos+curReadLen])
			if err != nil {
				pw.CloseWithError(err)
				return
			}

			totalReadLen -= curReadLen
			readPos += curReadLen
		}
		pw.Close()
	}()

	return pr, nil
}