tstor/src/sources/torrent/controller.go

package torrent

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

	"git.kmsign.ru/royalcat/tstor/pkg/kvsingle"
	"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, torrentFileProperties kv.Store[string, FileProperties], storage TorrentFileDeleter) *Controller {
	return &Controller{
		t:              t,
		storage:        storage,
		fileProperties: torrentFileProperties,
		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) ([]*FileController, error) {
	ctx, span := tracer.Start(ctx, "Files")
	defer span.End()

	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 := make([]*FileController, 0)
	for _, v := range s.t.Files() {
		if strings.Contains(v.Path(), "/.pad/") {
			continue
		}

		props := kvsingle.New(s.fileProperties, v.Path())
		ctl := NewFileController(v, props)
		files = append(files, ctl)
	}

	return files, nil
}

func (s *Controller) GetFile(ctx context.Context, file string) (*FileController, error) {
	files, err := s.Files(ctx)
	if err != nil {
		return nil, err
	}

	for _, v := range files {
		if v.Path() == file {
			return v, nil
		}
	}

	return nil, 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")

	err := s.fileProperties.Edit(ctx, f.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) {
		v.Excluded = true
		return v, nil
	})
	if err == kv.ErrKeyNotFound {
		err := s.fileProperties.Set(ctx, f.Path(), FileProperties{Excluded: true})
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}

	return s.storage.DeleteFile(f)
}

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) SetPriority(ctx context.Context, priority types.PiecePriority) error {
	// log := c.log.With(slog.Int("priority", int(priority)))

	for _, f := range c.t.Files() {
		err := c.setFilePriority(ctx, f, priority)
		if err != nil {
			return err
		}
	}

	return nil
}

const defaultPriority = types.PiecePriorityNone

func (c *Controller) Priority(ctx context.Context) (types.PiecePriority, error) {
	prio := defaultPriority
	files, err := c.Files(ctx)
	if err != nil {
		return 0, err
	}
	for _, v := range files {
		props, err := v.Properties(ctx)
		if err != nil {
			return 0, err
		}
		if props.Priority > prio {
			prio = props.Priority
		}
	}

	return prio, nil
}
func (c *Controller) setFilePriority(ctx context.Context, file *torrent.File, priority types.PiecePriority) error {
	err := c.fileProperties.Edit(ctx, file.Path(), func(ctx context.Context, v FileProperties) (FileProperties, error) {
		v.Priority = priority
		return v, nil
	})

	if err == kv.ErrKeyNotFound {
		seterr := c.fileProperties.Set(ctx, file.Path(), FileProperties{Priority: priority})
		if seterr != nil {
			return seterr
		}
		err = nil
	}

	if err != nil {
		return err
	}
	file.SetPriority(priority)
	return nil
}

func (c *Controller) initializeTorrentPriories(ctx context.Context) error {
	ctx, span := tracer.Start(ctx, "initializeTorrentPriories")
	defer span.End()
	log := c.log

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

	for _, file := range files {
		props, err := file.Properties(ctx)
		if err != nil {
			log.Error(ctx, "failed to get file properties", rlog.Error(err))
			continue
		}
		log = log.With(slog.Int("priority", int(props.Priority)))

		file.file.SetPriority(props.Priority)
	}

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

	return nil
}