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

package downloader

import (
	"fmt"
	"io"

	lru "github.com/hashicorp/golang-lru/v2"
	"gitlink.org.cn/cloudream/common/pkgs/iterator"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/db"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/downloader/strategy"
	"gitlink.org.cn/cloudream/jcs-pub/client/types"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/connectivity"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/pool"
)

const (
	DefaultMaxStripCacheCount = 128
)

type DownloadIterator = iterator.Iterator[*Downloading]

type DownloadReqeust struct {
	ObjectID types.ObjectID
	Offset   int64
	Length   int64
}

type downloadReqeust2 struct {
	Detail *types.ObjectDetail
	Raw    DownloadReqeust
}

type Downloading struct {
	Object  *types.Object
	File    io.ReadCloser // 文件流，如果文件不存在，那么为nil
	Request DownloadReqeust
}

type Downloader struct {
	strips   *StripCache
	cfg      Config
	conn     *connectivity.Collector
	stgPool  *pool.Pool
	selector *strategy.Selector
	db       *db.DB
}

func NewDownloader(cfg Config, conn *connectivity.Collector, stgPool *pool.Pool, sel *strategy.Selector, db *db.DB) *Downloader {
	if cfg.MaxStripCacheCount == 0 {
		cfg.MaxStripCacheCount = DefaultMaxStripCacheCount
	}

	ch, _ := lru.New[ECStripKey, ObjectECStrip](cfg.MaxStripCacheCount)
	return &Downloader{
		strips:   ch,
		cfg:      cfg,
		conn:     conn,
		stgPool:  stgPool,
		selector: sel,
		db:       db,
	}
}

func (d *Downloader) DownloadObjects(reqs []DownloadReqeust) DownloadIterator {
	objIDs := make([]types.ObjectID, len(reqs))
	for i, req := range reqs {
		objIDs[i] = req.ObjectID
	}

	if len(objIDs) == 0 {
		return iterator.Empty[*Downloading]()
	}

	objDetails, err := db.DoTx11(d.db, d.db.Object().BatchGetDetails, objIDs)
	if err != nil {
		return iterator.FuseError[*Downloading](fmt.Errorf("request to db: %w", err))
	}

	detailsMap := make(map[types.ObjectID]*types.ObjectDetail)
	for _, detail := range objDetails {
		d := detail
		detailsMap[detail.Object.ObjectID] = &d
	}

	req2s := make([]downloadReqeust2, len(reqs))
	for i, req := range reqs {
		req2s[i] = downloadReqeust2{
			Detail: detailsMap[req.ObjectID],
			Raw:    req,
		}
	}

	return NewDownloadObjectIterator(d, req2s)
}

func (d *Downloader) DownloadObjectByDetail(detail types.ObjectDetail, off int64, length int64) (*Downloading, error) {
	req2s := []downloadReqeust2{{
		Detail: &detail,
		Raw: DownloadReqeust{
			ObjectID: detail.Object.ObjectID,
			Offset:   off,
			Length:   length,
		},
	}}

	iter := NewDownloadObjectIterator(d, req2s)
	return iter.MoveNext()
}

func (d *Downloader) DownloadPackage(pkgID types.PackageID) DownloadIterator {
	details, err := db.DoTx11(d.db, d.db.Object().GetPackageObjectDetails, pkgID)
	if err != nil {
		return iterator.FuseError[*Downloading](fmt.Errorf("get package object details: %w", err))
	}

	req2s := make([]downloadReqeust2, len(details))
	for i, objDetail := range details {
		dt := objDetail
		req2s[i] = downloadReqeust2{
			Detail: &dt,
			Raw: DownloadReqeust{
				ObjectID: objDetail.Object.ObjectID,
				Offset:   0,
				Length:   objDetail.Object.Size,
			},
		}
	}

	return NewDownloadObjectIterator(d, req2s)
}

type ObjectECStrip struct {
	Data           []byte
	ObjectFileHash types.FileHash // 添加这条缓存时，Object的FileHash
}

type ECStripKey struct {
	ObjectID   types.ObjectID
	StripIndex int64
}

type StripCache = lru.Cache[ECStripKey, ObjectECStrip]