tstor/src/sources/torrent/stats.go

package torrent

import (
	"context"
	"encoding/json"
	"path"
	"slices"
	"time"

	"git.kmsign.ru/royalcat/tstor/src/log"
	"github.com/anacrolix/torrent/types/infohash"
	"github.com/dgraph-io/badger/v4"
)

func newStatsStore(metaDir string, lifetime time.Duration) (*statsStore, error) {
	db, err := badger.OpenManaged(
		badger.
			DefaultOptions(path.Join(metaDir, "stats")).
			WithNumVersionsToKeep(int(^uint(0) >> 1)).
			WithLogger(log.BadgerLogger("stats")), // Infinity
	)
	if err != nil {
		return nil, err
	}

	go func() {
		for n := range time.NewTimer(lifetime / 2).C {
			db.SetDiscardTs(uint64(n.Add(-lifetime).Unix()))
		}
	}()
	return &statsStore{
		db: db,
	}, nil
}

type statsStore struct {
	db *badger.DB
}

type TorrentStats struct {
	Timestamp        time.Time
	DownloadedBytes  uint64
	UploadedBytes    uint64
	TotalPeers       uint16
	ActivePeers      uint16
	ConnectedSeeders uint16
}

func (s TorrentStats) Same(o TorrentStats) bool {
	return s.DownloadedBytes == o.DownloadedBytes &&
		s.UploadedBytes == o.UploadedBytes &&
		s.TotalPeers == o.TotalPeers &&
		s.ActivePeers == o.ActivePeers &&
		s.ConnectedSeeders == o.ConnectedSeeders
}

func (r *statsStore) addStats(key []byte, stat TorrentStats) error {
	ts := uint64(stat.Timestamp.Unix())

	txn := r.db.NewTransactionAt(ts, true)
	defer txn.Discard()

	item, err := txn.Get(key)
	if err != nil && err != badger.ErrKeyNotFound {
		return err
	}

	if err != badger.ErrKeyNotFound {
		var prevStats TorrentStats
		err = item.Value(func(val []byte) error {
			return json.Unmarshal(val, &prevStats)
		})
		if err != nil {
			return err
		}

		if prevStats.Same(stat) {
			return nil
		}
	}

	data, err := json.Marshal(stat)
	if err != nil {
		return err
	}
	err = txn.Set(key, data)
	if err != nil {
		return err
	}

	return txn.CommitAt(ts, nil)
}

func (r *statsStore) AddTorrentStats(ih infohash.T, stat TorrentStats) error {
	return r.addStats(ih.Bytes(), stat)
}

const totalKey = "total"

func (r *statsStore) AddTotalStats(stat TorrentStats) error {
	return r.addStats([]byte(totalKey), stat)
}

func (r *statsStore) ReadTotalStatsHistory(ctx context.Context, since time.Time) ([]TorrentStats, error) {
	stats := []TorrentStats{}

	err := r.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.AllVersions = true
		opts.SinceTs = uint64(since.Unix())

		it := txn.NewKeyIterator([]byte(totalKey), opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			var stat TorrentStats
			err := item.Value(func(v []byte) error {
				return json.Unmarshal(v, &stat)
			})
			if err != nil {
				return err
			}

			stats = append(stats, stat)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	slices.SortFunc(stats, func(a, b TorrentStats) int {
		return a.Timestamp.Compare(b.Timestamp)
	})
	stats = slices.Compact(stats)
	return stats, nil
}

func (r *statsStore) ReadTorrentStatsHistory(ctx context.Context, since time.Time, ih infohash.T) ([]TorrentStats, error) {
	stats := []TorrentStats{}

	err := r.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.AllVersions = true
		opts.SinceTs = uint64(since.Unix())

		it := txn.NewKeyIterator(ih.Bytes(), opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			var stat TorrentStats
			err := item.Value(func(v []byte) error {
				return json.Unmarshal(v, &stat)
			})
			if err != nil {
				return err
			}

			stats = append(stats, stat)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	slices.SortFunc(stats, func(a, b TorrentStats) int {
		return a.Timestamp.Compare(b.Timestamp)
	})
	stats = slices.Compact(stats)
	return stats, nil
}

func (r *statsStore) ReadStatsHistory(ctx context.Context, since time.Time) ([]TorrentStats, error) {
	stats := []TorrentStats{}

	err := r.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.AllVersions = true
		opts.SinceTs = uint64(since.Unix())

		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			var stat TorrentStats
			err := item.Value(func(v []byte) error {
				return json.Unmarshal(v, &stat)
			})
			if err != nil {
				return err
			}

			stats = append(stats, stat)
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	slices.SortFunc(stats, func(a, b TorrentStats) int {
		return a.Timestamp.Compare(b.Timestamp)
	})
	stats = slices.Compact(stats)
	return stats, nil
}