137 lines
3.1 KiB
Go
137 lines
3.1 KiB
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())
|
|
}
|