tstor/daemons/torrent/piece_completion.go

package torrent

import (
	"context"
	"encoding/binary"
	"fmt"

	"git.kmsign.ru/royalcat/tstor/src/logwrap"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
	"github.com/royalcat/kv"
	"github.com/royalcat/kv/kvbadger"
)

type PieceCompletionState byte

const (
	PieceNotComplete PieceCompletionState = 0
	PieceComplete    PieceCompletionState = 1<<8 - 1
)

var _ kv.Binary = (*PieceCompletionState)(nil)

// MarshalBinary implements kv.Binary.
func (p PieceCompletionState) MarshalBinary() (data []byte, err error) {
	return []byte{byte(p)}, nil
}

// UnmarshalBinary implements kv.Binary.
func (p *PieceCompletionState) UnmarshalBinary(data []byte) error {
	if len(data) != 1 {
		return fmt.Errorf("bad length")
	}

	switch PieceCompletionState(data[0]) {
	case PieceComplete:
		*p = PieceComplete
	case PieceNotComplete:
		*p = PieceNotComplete
	default:
		*p = PieceNotComplete
	}

	return nil
}

func pieceCompletionState(i bool) PieceCompletionState {
	if i {
		return PieceComplete
	}
	return PieceNotComplete
}

type pieceKey metainfo.PieceKey

const pieceKeySize = metainfo.HashSize + 4

var _ kv.Binary = (*pieceKey)(nil)

// const delimeter rune = 0x1F

// MarshalBinary implements kv.Binary.
func (pk pieceKey) MarshalBinary() (data []byte, err error) {
	key := make([]byte, 0, pieceKeySize)
	key = append(key, pk.InfoHash.Bytes()...)
	key = binary.BigEndian.AppendUint32(key, uint32(pk.Index))
	return key, nil
}

// UnmarshalBinary implements kv.Binary.
func (p *pieceKey) UnmarshalBinary(data []byte) error {
	if len(data) < pieceKeySize {
		return fmt.Errorf("data too short")
	}
	p.InfoHash = metainfo.Hash(data[:metainfo.HashSize])
	p.Index = int(binary.BigEndian.Uint32(data[metainfo.HashSize:]))
	return nil
}

type badgerPieceCompletion struct {
	db kv.Store[pieceKey, PieceCompletionState]
}

var _ storage.PieceCompletion = (*badgerPieceCompletion)(nil)

func newPieceCompletion(dir string) (storage.PieceCompletion, error) {
	opts := kvbadger.DefaultOptions[PieceCompletionState](dir)
	opts.Codec = kv.CodecBinary[PieceCompletionState, *PieceCompletionState]{}
	opts.BadgerOptions = opts.BadgerOptions.WithLogger(logwrap.BadgerLogger("torrent-client", "piece-completion"))

	db, err := kvbadger.NewBagerKVBinaryKey[pieceKey, PieceCompletionState](opts)
	if err != nil {
		return nil, err
	}

	return &badgerPieceCompletion{
		db: db,
	}, nil
}

func (c *badgerPieceCompletion) Get(pk metainfo.PieceKey) (completion storage.Completion, err error) {
	ctx := context.Background()

	state, err := c.db.Get(ctx, pieceKey(pk))
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return completion, nil
		}
		return completion, err
	}

	if state == PieceComplete {
		return storage.Completion{
			Complete: true,
			Ok:       true,
		}, nil
	}

	return storage.Completion{
		Complete: false,
		Ok:       true,
	}, nil
}

func (me badgerPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
	ctx := context.Background()

	if c, err := me.Get(pk); err == nil && c.Ok && c.Complete == b {
		return nil
	}

	return me.db.Set(ctx, pieceKey(pk), pieceCompletionState(b))
}

func (me *badgerPieceCompletion) Close() error {
	return me.db.Close(context.Background())
}