package qbittorrent import ( "context" "fmt" "slices" "time" "git.kmsign.ru/royalcat/tstor/pkg/qbittorrent" "github.com/hashicorp/golang-lru/v2/expirable" "github.com/royalcat/btrgo/btrsync" ) type cacheClient struct { qb qbittorrent.Client propertiesCache *expirable.LRU[string, qbittorrent.TorrentProperties] torrentsCache *expirable.LRU[string, qbittorrent.TorrentInfo] pieceCache btrsync.MapOf[pieceKey, int] } type pieceKey struct { hash string index int } func wrapClient(qb qbittorrent.Client) *cacheClient { const ( cacheSize = 5000 cacheTTL = time.Minute ) return &cacheClient{ qb: qb, propertiesCache: expirable.NewLRU[string, qbittorrent.TorrentProperties](cacheSize, nil, cacheTTL), torrentsCache: expirable.NewLRU[string, qbittorrent.TorrentInfo](cacheSize, nil, cacheTTL), pieceCache: btrsync.MapOf[pieceKey, int]{}, } } func (f *cacheClient) getInfo(ctx context.Context, hash string) (*qbittorrent.TorrentInfo, error) { if v, ok := f.torrentsCache.Get(hash); ok { return &v, nil } infos, err := f.qb.Torrent().GetTorrents(ctx, &qbittorrent.TorrentOption{ Hashes: []string{hash}, }) if err != nil { return nil, fmt.Errorf("error to check torrent existence: %w", err) } if len(infos) == 0 { return nil, nil } if len(infos) > 1 { return nil, fmt.Errorf("multiple torrents with the same hash") } f.torrentsCache.Add(hash, *infos[0]) return infos[0], nil } func (f *cacheClient) getProperties(ctx context.Context, hash string) (*qbittorrent.TorrentProperties, error) { if v, ok := f.propertiesCache.Get(hash); ok { return &v, nil } info, err := f.qb.Torrent().GetProperties(ctx, hash) if err != nil { return nil, err } f.propertiesCache.Add(hash, *info) return info, nil } func (f *cacheClient) listContent(ctx context.Context, hash string) ([]*qbittorrent.TorrentContent, error) { contents, err := f.qb.Torrent().GetContents(ctx, hash) if err != nil { return nil, err } return contents, nil } func (f *cacheClient) getContent(ctx context.Context, hash string, contentIndex int) (*qbittorrent.TorrentContent, error) { contents, err := f.qb.Torrent().GetContents(ctx, hash, contentIndex) if err != nil { return nil, err } contentI := slices.IndexFunc(contents, func(c *qbittorrent.TorrentContent) bool { return c.Index == contentIndex }) if contentI == -1 { return nil, fmt.Errorf("content not found") } return contents[contentI], nil } func (f *cacheClient) isPieceComplete(ctx context.Context, hash string, pieceIndex int) (bool, error) { cachedPieceState, ok := f.pieceCache.Load(pieceKey{hash: hash, index: pieceIndex}) if ok && cachedPieceState == 2 { return true, nil } completion, err := f.qb.Torrent().GetPiecesStates(ctx, hash) if err != nil { return false, err } for i, v := range completion { f.pieceCache.Store(pieceKey{hash: hash, index: i}, v) } if completion[pieceIndex] == 2 { return true, nil } return false, nil } func (f *cacheClient) waitPieceToComplete(ctx context.Context, hash string, pieceIndex int) error { const checkingInterval = 1 * time.Second ok, err := f.isPieceComplete(ctx, hash, pieceIndex) if err != nil { return err } if ok { return nil } if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) < checkingInterval { return context.DeadlineExceeded } ticker := time.NewTicker(checkingInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return ctx.Err() case <-ticker.C: ok, err := f.isPieceComplete(ctx, hash, pieceIndex) if err != nil { return err } if ok { return nil } if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) < checkingInterval { return context.DeadlineExceeded } } } }