tstor/daemons/torrent/infobytes.go

package torrent

import (
	"bytes"
	"errors"
	"fmt"
	"path/filepath"

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

var errNotFound = errors.New("not found")

type infoBytesStore struct {
	db *badger.DB
}

func newInfoBytesStore(metaDir string) (*infoBytesStore, error) {
	opts := badger.
		DefaultOptions(filepath.Join(metaDir, "infobytes")).
		WithLogger(logwrap.BadgerLogger("torrent-client", "infobytes"))
	db, err := badger.Open(opts)
	if err != nil {
		return nil, err
	}
	return &infoBytesStore{db}, nil
}

func (k *infoBytesStore) GetBytes(ih infohash.T) ([]byte, error) {
	var data []byte
	err := k.db.View(func(tx *badger.Txn) error {
		item, err := tx.Get(ih.Bytes())
		if err != nil {
			if err == badger.ErrKeyNotFound {
				return errNotFound
			}

			return fmt.Errorf("error getting value: %w", err)
		}

		data, err = item.ValueCopy(data)
		return err
	})
	return data, err
}

func (k *infoBytesStore) Get(ih infohash.T) (*metainfo.MetaInfo, error) {
	data, err := k.GetBytes(ih)
	if err != nil {
		return nil, err
	}

	return metainfo.Load(bytes.NewReader(data))
}

func (me *infoBytesStore) SetBytes(ih infohash.T, data []byte) error {
	return me.db.Update(func(txn *badger.Txn) error {
		item, err := txn.Get(ih.Bytes())
		if err != nil {
			if err == badger.ErrKeyNotFound {
				return txn.Set(ih.Bytes(), data)
			}
			return err
		}

		return item.Value(func(val []byte) error {
			if !bytes.Equal(val, data) {
				return txn.Set(ih.Bytes(), data)
			}
			return nil
		})
	})
}

func (me *infoBytesStore) Set(ih infohash.T, info metainfo.MetaInfo) error {
	return me.SetBytes(ih, info.InfoBytes)
}

func (k *infoBytesStore) Delete(ih infohash.T) error {
	return k.db.Update(func(txn *badger.Txn) error {
		return txn.Delete(ih.Bytes())
	})
}

func (me *infoBytesStore) Close() error {
	return me.db.Close()
}