tstor/src/sources/torrent/dht_fileitem_store.go

package torrent

import (
	"bytes"
	"encoding/gob"
	"time"

	dlog "git.kmsign.ru/royalcat/tstor/src/log"
	"github.com/anacrolix/dht/v2/bep44"
	"github.com/dgraph-io/badger/v4"
)

var _ bep44.Store = &dhtFileItemStore{}

type dhtFileItemStore struct {
	ttl time.Duration
	db  *badger.DB
}

func newDHTStore(path string, itemsTTL time.Duration) (*dhtFileItemStore, error) {
	opts := badger.DefaultOptions(path).
		WithLogger(dlog.BadgerLogger("torrent-client", "dht-item-store")).
		WithValueLogFileSize(1<<26 - 1)

	db, err := badger.Open(opts)
	if err != nil {
		return nil, err
	}

	err = db.RunValueLogGC(0.5)
	if err != nil && err != badger.ErrNoRewrite {
		return nil, err
	}

	return &dhtFileItemStore{
		db:  db,
		ttl: itemsTTL,
	}, nil
}

func (fis *dhtFileItemStore) Put(i *bep44.Item) error {
	tx := fis.db.NewTransaction(true)
	defer tx.Discard()

	key := i.Target()
	var value bytes.Buffer

	enc := gob.NewEncoder(&value)
	if err := enc.Encode(i); err != nil {
		return err
	}

	e := badger.NewEntry(key[:], value.Bytes()).WithTTL(fis.ttl)
	if err := tx.SetEntry(e); err != nil {
		return err
	}

	return tx.Commit()
}

func (fis *dhtFileItemStore) Get(t bep44.Target) (*bep44.Item, error) {
	tx := fis.db.NewTransaction(false)
	defer tx.Discard()

	dbi, err := tx.Get(t[:])
	if err == badger.ErrKeyNotFound {
		return nil, bep44.ErrItemNotFound
	}
	if err != nil {
		return nil, err
	}
	valb, err := dbi.ValueCopy(nil)
	if err != nil {
		return nil, err
	}

	buf := bytes.NewBuffer(valb)
	dec := gob.NewDecoder(buf)
	var i *bep44.Item
	if err := dec.Decode(&i); err != nil {
		return nil, err
	}

	return i, nil
}

func (fis *dhtFileItemStore) Del(t bep44.Target) error {
	tx := fis.db.NewTransaction(true)
	defer tx.Discard()

	err := tx.Delete(t[:])
	if err == badger.ErrKeyNotFound {
		return nil
	}
	if err != nil {
		return err
	}

	err = tx.Commit()
	if err == badger.ErrKeyNotFound {
		return nil
	}
	if err != nil {
		return err
	}

	return nil
}

func (fis *dhtFileItemStore) Close() error {
	return fis.db.Close()
}