tstor/daemons/qbittorrent/client.go

package qbittorrent

import (
	"context"
	"fmt"
	"slices"
	"time"

	"git.kmsign.ru/royalcat/tstor/pkg/qbittorrent"
	"github.com/creativecreature/sturdyc"
	"github.com/hashicorp/golang-lru/v2/expirable"
	"github.com/royalcat/btrgo/btrsync"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/metric"
)

var meter = otel.Meter("git.kmsign.ru/royalcat/tstor/daemons/qbittorrent")

type cacheClient struct {
	qb qbittorrent.Client

	propertiesCache *expirable.LRU[string, qbittorrent.TorrentProperties]

	pieceCache btrsync.MapOf[pieceKey, int]

	infoClient *sturdyc.Client[*qbittorrent.TorrentInfo]
}

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),

		infoClient: sturdyc.New[*qbittorrent.TorrentInfo](cacheSize, 1, cacheTTL, 10,
			sturdyc.WithEarlyRefreshes(time.Minute, time.Minute*5, time.Second*10),
			sturdyc.WithRefreshCoalescing(100, time.Second/4),
			sturdyc.WithMetrics(newSturdycMetrics()),
		),

		pieceCache: btrsync.MapOf[pieceKey, int]{},
	}
}

func (f *cacheClient) getInfo(ctx context.Context, hash string) (*qbittorrent.TorrentInfo, error) {
	out, err := f.infoClient.GetOrFetchBatch(ctx, []string{hash},
		f.infoClient.BatchKeyFn(""),
		func(ctx context.Context, ids []string) (map[string]*qbittorrent.TorrentInfo, error) {
			infos, err := f.qb.Torrent().GetTorrents(ctx, &qbittorrent.TorrentOption{
				Hashes: ids,
			})
			if err != nil {
				return nil, fmt.Errorf("error to get torrents: %w", err)
			}

			out := make(map[string]*qbittorrent.TorrentInfo)
			for _, info := range infos {
				out[info.Hash] = info
			}

			return out, nil
		},
	)
	if err != nil {
		return nil, err
	}

	if out[hash] == nil {
		return nil, nil
	}

	return out[hash], 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
			}
		}
	}
}

type sturdycMetrics struct {
	ctx context.Context

	cacheHit       metric.Int64Counter
	cacheMiss      metric.Int64Counter
	refresh        metric.Int64Counter
	missing        metric.Int64Counter
	forcedEviction metric.Int64Counter
	entryEviction  metric.Int64Counter
	batchSize      metric.Int64Histogram

	observeCacheSize func() int
}

var _ sturdyc.MetricsRecorder = (*sturdycMetrics)(nil)

func newSturdycMetrics() *sturdycMetrics {
	m := &sturdycMetrics{
		ctx:            context.Background(),
		cacheHit:       must(meter.Int64Counter("sturdyc_cache_hit")),
		cacheMiss:      must(meter.Int64Counter("sturdyc_cache_miss")),
		refresh:        must(meter.Int64Counter("sturdyc_cache_refresh")),
		missing:        must(meter.Int64Counter("sturdyc_cache_missing")),
		forcedEviction: must(meter.Int64Counter("sturdyc_cache_forced_eviction")),
		entryEviction:  must(meter.Int64Counter("sturdyc_cache_entry_eviction")),
		batchSize:      must(meter.Int64Histogram("sturdyc_cache_batch_size")),
	}

	must(meter.Int64ObservableGauge("sturdyc_cache_size",
		metric.WithInt64Callback(func(ctx context.Context, io metric.Int64Observer) error {
			if m.observeCacheSize == nil {
				return nil
			}
			io.Observe(int64(m.observeCacheSize()))
			return nil
		})))

	return m
}

func (s *sturdycMetrics) CacheHit() {
	s.cacheHit.Add(s.ctx, 1)
}

func (s *sturdycMetrics) CacheMiss() {
	s.cacheMiss.Add(s.ctx, 1)
}

func (s *sturdycMetrics) Refresh() {
	s.refresh.Add(s.ctx, 1)
}

func (s *sturdycMetrics) MissingRecord() {
	s.missing.Add(s.ctx, 1)
}

func (s *sturdycMetrics) ForcedEviction() {
	s.forcedEviction.Add(s.ctx, 1)
}

func (s *sturdycMetrics) CacheBatchRefreshSize(size int) {
	s.batchSize.Record(s.ctx, int64(size))
}

func (s *sturdycMetrics) ObserveCacheSize(callback func() int) {
	s.observeCacheSize = callback
}

func (s *sturdycMetrics) EntriesEvicted(evictd int) {
	s.entryEviction.Add(s.ctx, int64(evictd))
}

func (s *sturdycMetrics) ShardIndex(int) {
	return
}

func must[T any](v T, err error) T {
	if err != nil {
		panic(err)
	}
	return v
}