tstor/src/sources/torrent/controller.go

package torrent

import (
	"context"
	"log/slog"
	"slices"
	"strings"

	"git.kmsign.ru/royalcat/tstor/pkg/rlog"
	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/types"
	"github.com/royalcat/kv"
)

type TorrentFileDeleter interface {
	DeleteFile(file *torrent.File) error
}

type FileProperties struct {
	Excluded bool                `json:"excluded"`
	Priority types.PiecePriority `json:"priority"`
}

type Controller struct {
	torrentFilePath string
	t               *torrent.Torrent
	storage         TorrentFileDeleter
	fileProperties  kv.Store[string, FileProperties]
	log             *rlog.Logger
}

func newController(t *torrent.Torrent, fileProperties kv.Store[string, FileProperties], storage TorrentFileDeleter) *Controller {
	return &Controller{
		t:              t,
		storage:        storage,
		fileProperties: fileProperties,
		log:            rlog.Component("torrent-client", "controller").With(slog.String("infohash", t.InfoHash().HexString())),
	}
}

func (s *Controller) TorrentFilePath() string {
	return s.torrentFilePath
}

func (s *Controller) Torrent() *torrent.Torrent {
	return s.t
}

func (c *Controller) Name() string {
	<-c.t.GotInfo()
	if name := c.t.Name(); name != "" {
		return name
	}

	return c.InfoHash()
}

func (s *Controller) InfoHash() string {
	<-s.t.GotInfo()
	return s.t.InfoHash().HexString()
}

func (s *Controller) BytesCompleted() int64 {
	<-s.t.GotInfo()
	return s.t.BytesCompleted()
}

func (s *Controller) BytesMissing() int64 {
	<-s.t.GotInfo()
	return s.t.BytesMissing()
}

func (s *Controller) Length() int64 {
	<-s.t.GotInfo()
	return s.t.Length()
}

func (s *Controller) Files(ctx context.Context) ([]*torrent.File, error) {
	fps := map[string]FileProperties{}
	err := s.fileProperties.Range(ctx, func(k string, v FileProperties) error {
		fps[k] = v
		return nil
	})
	if err != nil {
		return nil, err
	}

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-s.t.GotInfo():
	}

	files := s.t.Files()
	files = slices.DeleteFunc(files, func(file *torrent.File) bool {
		if file == nil {
			return true
		}

		p := file.Path()
		if strings.Contains(p, "/.pad/") {
			return true
		}
		if props, ok := fps[p]; ok && props.Excluded {
			return true
		}

		return false
	})

	return files, nil
}

func Map[T, U any](ts []T, f func(T) U) []U {
	us := make([]U, len(ts))
	for i := range ts {
		us[i] = f(ts[i])
	}
	return us
}

func (s *Controller) ExcludeFile(ctx context.Context, f *torrent.File) error {
	log := s.log.With(slog.String("file", f.Path()))
	log.Info(ctx, "excluding file")

	return s.fileProperties.Edit(ctx, f.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) {
		v.Excluded = true
		return v, nil
	})
}

func (s *Controller) isFileComplete(startIndex int, endIndex int) bool {
	for i := startIndex; i < endIndex; i++ {
		if !s.t.Piece(i).State().Complete {
			return false
		}
	}
	return true
}

func (s *Controller) ValidateTorrent(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-s.t.GotInfo():
	}

	for i := 0; i < s.t.NumPieces(); i++ {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		s.t.Piece(i).VerifyData()
	}

	return nil
}

func (c *Controller) SetFilePriority(ctx context.Context, file *torrent.File, priority types.PiecePriority) error {
	log := c.log.With(slog.String("file", file.Path()), slog.Int("priority", int(priority)))
	log.Info(ctx, "set pritority for file")

	err := c.fileProperties.Edit(ctx, file.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) {
		v.Priority = priority
		return v, nil
	})
	if err != nil {
		return err
	}

	file.SetPriority(priority)
	return nil
}

func (c *Controller) initializeTorrentPriories(ctx context.Context) error {
	log := c.log.WithComponent("initializeTorrentPriories")

	files, err := c.Files(ctx)
	if err != nil {
		return err
	}

	for _, file := range files {
		if file == nil {
			continue
		}

		props, err := c.fileProperties.Get(ctx, file.Path())
		if err != nil {
			if err == kv.ErrKeyNotFound {
				continue
			}
			log.Error(ctx, "failed to get file priority", rlog.Error(err))
		}

		file.SetPriority(props.Priority)

	}

	log.Info(ctx, "torrent initialization complete", slog.String("infohash", c.InfoHash()), slog.String("torrent_name", c.Name()))

	return nil
}