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/internal/speedstats"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/connectivity"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/pool"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/sysevent"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

const (
	DefaultMaxStripCacheCount = 128
)

type DownloadIterator = iterator.Iterator[*Downloading]

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

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

type Downloading struct {
	Object  *jcstypes.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
	speedStats *speedstats.SpeedStats
	evtPub     *sysevent.Publisher
}

func NewDownloader(cfg Config, conn *connectivity.Collector, stgPool *pool.Pool, sel *strategy.Selector, db *db.DB, speedStats *speedstats.SpeedStats, evtPub *sysevent.Publisher) *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,
		speedStats: speedStats,
		evtPub:     evtPub,
	}
}

func (d *Downloader) DownloadObjects(reqs []DownloadReqeust) DownloadIterator {
	objIDs := make([]jcstypes.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[jcstypes.ObjectID]*jcstypes.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 jcstypes.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 jcstypes.PackageID, prefix string) (jcstypes.Package, DownloadIterator, error) {
	pkg, details, err := db.DoTx02(d.db, func(tx db.SQLContext) (jcstypes.Package, []jcstypes.ObjectDetail, error) {
		pkg, err := d.db.Package().GetByID(tx, pkgID)
		if err != nil {
			return jcstypes.Package{}, nil, err
		}

		var details []jcstypes.ObjectDetail
		if prefix != "" {
			objs, err := d.db.Object().GetWithPathPrefix(tx, pkgID, prefix)
			if err != nil {
				return jcstypes.Package{}, nil, err
			}

			objIDs := make([]jcstypes.ObjectID, len(objs))
			for i, obj := range objs {
				objIDs[i] = obj.ObjectID
			}

			allBlocks, err := d.db.ObjectBlock().BatchGetByObjectID(tx, objIDs)
			if err != nil {
				return jcstypes.Package{}, nil, err
			}

			allPinnedObjs, err := d.db.PinnedObject().BatchGetByObjectID(tx, objIDs)
			if err != nil {
				return jcstypes.Package{}, nil, err

			}
			details = make([]jcstypes.ObjectDetail, 0, len(objs))
			for _, obj := range objs {
				detail := jcstypes.ObjectDetail{
					Object: obj,
				}
				details = append(details, detail)
			}

			jcstypes.DetailsFillObjectBlocks(details, allBlocks)
			jcstypes.DetailsFillPinnedAt(details, allPinnedObjs)
		} else {
			details, err = d.db.Object().GetPackageObjectDetails(tx, pkgID)
			if err != nil {
				return jcstypes.Package{}, nil, err
			}
		}

		return pkg, details, nil
	})
	if err != nil {
		return jcstypes.Package{}, nil, 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 pkg, NewDownloadObjectIterator(d, req2s), nil
}

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

type ECStripKey struct {
	ObjectID   jcstypes.ObjectID
	StripIndex int64
}

type StripCache = lru.Cache[ECStripKey, ObjectECStrip]