tstor/src/host/vfs/torrent.go

package vfs

import (
	"context"
	"io"
	"io/fs"
	"path"
	"slices"
	"strings"
	"sync"
	"time"

	"git.kmsign.ru/royalcat/tstor/src/host/controller"
	"github.com/anacrolix/torrent"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"golang.org/x/exp/maps"
)

type TorrentFs struct {
	name string

	mu      sync.Mutex
	Torrent *controller.Torrent

	filesCache map[string]File

	resolver *resolver
}

var _ Filesystem = (*TorrentFs)(nil)

func NewTorrentFs(name string, c *controller.Torrent) *TorrentFs {
	return &TorrentFs{
		name:     name,
		Torrent:  c,
		resolver: newResolver(ArchiveFactories),
	}
}

var _ fs.DirEntry = (*TorrentFs)(nil)

// Name implements fs.DirEntry.
func (tfs *TorrentFs) Name() string {
	return tfs.name
}

// Info implements fs.DirEntry.
func (tfs *TorrentFs) Info() (fs.FileInfo, error) {
	return tfs, nil
}

// IsDir implements fs.DirEntry.
func (tfs *TorrentFs) IsDir() bool {
	return true
}

// Type implements fs.DirEntry.
func (tfs *TorrentFs) Type() fs.FileMode {
	return fs.ModeDir
}

// ModTime implements fs.FileInfo.
func (tfs *TorrentFs) ModTime() time.Time {
	return time.Time{}
}

// Mode implements fs.FileInfo.
func (tfs *TorrentFs) Mode() fs.FileMode {
	return fs.ModeDir
}

// Size implements fs.FileInfo.
func (tfs *TorrentFs) Size() int64 {
	return 0
}

// Sys implements fs.FileInfo.
func (tfs *TorrentFs) Sys() any {
	return nil
}

// FsName implements Filesystem.
func (tfs *TorrentFs) FsName() string {
	return "torrentfs"
}

func (fs *TorrentFs) files(ctx context.Context) (map[string]File, error) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	if fs.filesCache != nil {
		return fs.filesCache, nil
	}

	ctx, span := tracer.Start(ctx, "files", fs.traceAttrs())
	defer span.End()

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

	fs.filesCache = make(map[string]File)
	for _, file := range files {
		file.SetPriority(torrent.PiecePriorityNormal)
		p := AbsPath(file.Path())
		tf, err := openTorrentFile(ctx, path.Base(p), file)
		if err != nil {
			return nil, err
		}
		fs.filesCache[p] = tf
	}

	// TODO optional
	if len(fs.filesCache) == 1 && fs.resolver.isNestedFs(fs.Torrent.Name()) {
		filepath := "/" + fs.Torrent.Name()
		if file, ok := fs.filesCache[filepath]; ok {
			nestedFs, err := fs.resolver.nestedFs(ctx, filepath, file)
			if err != nil {
				return nil, err
			}
			if nestedFs == nil {
				goto DEFAULT_DIR // FIXME
			}
			fs.filesCache, err = fs.listFilesRecursive(ctx, nestedFs, "/")
			if err != nil {
				return nil, err
			}

			return fs.filesCache, nil
		}

	}

DEFAULT_DIR:
	rootDir := "/" + fs.Torrent.Name() + "/"
	singleDir := true
	for k, _ := range fs.filesCache {
		if !strings.HasPrefix(k, rootDir) {
			singleDir = false
		}
	}
	if singleDir {
		for k, f := range fs.filesCache {
			delete(fs.filesCache, k)
			k, _ = strings.CutPrefix(k, rootDir)
			k = AbsPath(k)
			fs.filesCache[k] = f
		}
	}

	return fs.filesCache, nil
}

// func anyPeerHasFiles(file *torrent.File) bool {
// 	for _, conn := range file.Torrent().PeerConns() {
// 		if bitmapHaveFile(conn.PeerPieces(), file) {
// 			return true
// 		}
// 	}
// 	return false
// }

// func bitmapHaveFile(bitmap *roaring.Bitmap, file *torrent.File) bool {
// 	for i := file.BeginPieceIndex(); i < file.EndPieceIndex(); i++ {
// 		if !bitmap.ContainsInt(i) {
// 			return false
// 		}
// 	}
// 	return true
// }

func (fs *TorrentFs) listFilesRecursive(ctx context.Context, vfs Filesystem, start string) (map[string]File, error) {
	ctx, span := tracer.Start(ctx, "listFilesRecursive",
		fs.traceAttrs(attribute.String("start", start)),
	)
	defer span.End()

	out := make(map[string]File, 0)
	entries, err := vfs.ReadDir(ctx, start)
	if err != nil {
		return nil, err
	}
	for _, entry := range entries {
		filename := path.Join(start, entry.Name())
		if entry.IsDir() {
			rec, err := fs.listFilesRecursive(ctx, vfs, filename)
			if err != nil {
				return nil, err
			}
			maps.Copy(out, rec)
		} else {
			file, err := vfs.Open(ctx, filename)
			if err != nil {
				return nil, err
			}
			out[filename] = file
		}
	}

	return out, nil
}

func (fs *TorrentFs) rawOpen(ctx context.Context, filename string) (file File, err error) {
	ctx, span := tracer.Start(ctx, "rawOpen",
		fs.traceAttrs(attribute.String("filename", filename)),
	)
	defer func() {
		if err != nil {
			span.RecordError(err)
		}
		span.End()
	}()

	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}
	file, err = getFile(files, filename)
	return file, err
}

func (fs *TorrentFs) rawStat(ctx context.Context, filename string) (fs.FileInfo, error) {
	ctx, span := tracer.Start(ctx, "rawStat",
		fs.traceAttrs(attribute.String("filename", filename)),
	)
	defer span.End()

	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}
	file, err := getFile(files, filename)
	if err != nil {
		return nil, err
	}
	return file.Info()
}

func (fs *TorrentFs) traceAttrs(add ...attribute.KeyValue) trace.SpanStartOption {
	return trace.WithAttributes(append([]attribute.KeyValue{
		attribute.String("fs", fs.FsName()),
		attribute.String("torrent", fs.Torrent.Name()),
		attribute.String("infohash", fs.Torrent.InfoHash()),
	}, add...)...)
}

// Stat implements Filesystem.
func (fs *TorrentFs) Stat(ctx context.Context, filename string) (fs.FileInfo, error) {
	ctx, span := tracer.Start(ctx, "Stat",
		fs.traceAttrs(attribute.String("filename", filename)),
	)
	defer span.End()

	if isRoot(filename) {
		return fs, nil
	}

	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, filename, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.Stat(ctx, nestedFsPath)
	}

	return fs.rawStat(ctx, fsPath)
}

func (fs *TorrentFs) Open(ctx context.Context, filename string) (File, error) {
	ctx, span := tracer.Start(ctx, "Open",
		fs.traceAttrs(attribute.String("filename", filename)),
	)
	defer span.End()

	if isRoot(filename) {
		return newDirFile(fs.name), nil
	}

	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, filename, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.Open(ctx, nestedFsPath)
	}

	return fs.rawOpen(ctx, fsPath)
}

func (fs *TorrentFs) ReadDir(ctx context.Context, name string) ([]fs.DirEntry, error) {
	ctx, span := tracer.Start(ctx, "ReadDir",
		fs.traceAttrs(attribute.String("name", name)),
	)
	defer span.End()

	fsPath, nestedFs, nestedFsPath, err := fs.resolver.resolvePath(ctx, name, fs.rawOpen)
	if err != nil {
		return nil, err
	}
	if nestedFs != nil {
		return nestedFs.ReadDir(ctx, nestedFsPath)
	}
	files, err := fs.files(ctx)
	if err != nil {
		return nil, err
	}

	return listDirFromFiles(files, fsPath)
}

func (fs *TorrentFs) Unlink(ctx context.Context, name string) error {
	ctx, span := tracer.Start(ctx, "Unlink",
		fs.traceAttrs(attribute.String("name", name)),
	)
	defer span.End()

	name = AbsPath(name)

	fs.mu.Lock()
	defer fs.mu.Unlock()

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

	if !slices.Contains(maps.Keys(files), name) {
		return ErrNotExist
	}

	file := files[name]
	delete(fs.filesCache, name)

	tfile, ok := file.(*torrentFile)
	if !ok {
		return ErrNotImplemented
	}

	return fs.Torrent.ExcludeFile(ctx, tfile.file)
}

var _ File = (*torrentFile)(nil)

type torrentFile struct {
	name string

	mu sync.Mutex

	tr torrent.Reader

	lastReadTimeout time.Time

	file *torrent.File
}

const secondaryTimeout = time.Hour

func openTorrentFile(ctx context.Context, name string, file *torrent.File) (*torrentFile, error) {
	// select {
	// case <-file.Torrent().GotInfo():
	// 	break
	// case <-ctx.Done():
	// 	return nil, ctx.Err()
	// }

	r := file.NewReader()
	r.SetReadahead(1024 * 1024 * 16) // TODO configurable
	// r.SetResponsive()

	return &torrentFile{
		name: name,
		tr:   r,
		file: file,
	}, nil
}

// Name implements File.
func (tf *torrentFile) Name() string {
	return tf.name
}

// Type implements File.
func (tf *torrentFile) Type() fs.FileMode {
	return roMode | fs.ModeDir
}

func (tf *torrentFile) Info() (fs.FileInfo, error) {
	return newFileInfo(tf.name, tf.file.Length()), nil
}

func (tf *torrentFile) Size() int64 {
	return tf.file.Length()
}

func (tf *torrentFile) IsDir() bool {
	return false
}

func (rw *torrentFile) Close(ctx context.Context) error {
	rw.mu.Lock()
	defer rw.mu.Unlock()

	return rw.tr.Close()
}

// Read implements ctxio.Reader.
func (tf *torrentFile) Read(ctx context.Context, p []byte) (n int, err error) {
	ctx, span := tracer.Start(ctx, "Read",
		trace.WithAttributes(attribute.Int("length", len(p))),
	)
	defer func() {
		span.SetAttributes(attribute.Int("read", n))
		span.End()
	}()

	tf.mu.Lock()
	defer tf.mu.Unlock()

	if time.Since(tf.lastReadTimeout) < secondaryTimeout { // make short timeout for already faliled files
		span.SetAttributes(attribute.Bool("short_timeout", true))
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Millisecond)
		defer cancel()
	}
	defer func() {
		if err == context.DeadlineExceeded {
			tf.lastReadTimeout = time.Now()
		}
	}()

	return tf.tr.ReadContext(ctx, p)
}

func (tf *torrentFile) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) {
	ctx, span := tracer.Start(ctx, "ReadAt",
		trace.WithAttributes(attribute.Int("length", len(p)), attribute.Int64("offset", off)),
	)
	defer func() {
		span.SetAttributes(attribute.Int("read", n))
		span.End()
	}()

	tf.mu.Lock()
	defer tf.mu.Unlock()

	if time.Since(tf.lastReadTimeout) < secondaryTimeout { // make short timeout for already faliled files
		span.SetAttributes(attribute.Bool("short_timeout", true))
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Millisecond)
		defer cancel()
	}
	defer func() {
		if err == context.DeadlineExceeded {
			tf.lastReadTimeout = time.Now()
		}
	}()

	_, err = tf.tr.Seek(off, io.SeekStart)
	if err != nil {
		return 0, err
	}

	// return tf.tr.ReadContext(ctx, p)
	n, err = readAtLeast(ctx, tf.tr, p, len(p))

	_, err = tf.tr.Seek(0, io.SeekStart)
	if err != nil {
		return 0, err
	}

	return n, err
}

func readAtLeast(ctx context.Context, r torrent.Reader, buf []byte, min int) (n int, err error) {
	if len(buf) < min {
		return 0, io.ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int

		nn, err = r.ReadContext(ctx, buf[n:])
		n += nn
	}
	if n >= min {
		err = nil
	} else if n > 0 && err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}