github.com/anacrolix/torrent@v1.61.0/storage/bolt-piece.go (about)

     1  //go:build !noboltdb && !wasm
     2  // +build !noboltdb,!wasm
     3  
     4  package storage
     5  
     6  import (
     7  	"encoding/binary"
     8  	"io"
     9  
    10  	"go.etcd.io/bbolt"
    11  
    12  	"github.com/anacrolix/torrent/metainfo"
    13  )
    14  
    15  type boltPiece struct {
    16  	db  *bbolt.DB
    17  	p   metainfo.Piece
    18  	ih  metainfo.Hash
    19  	key [24]byte
    20  }
    21  
    22  var (
    23  	_             PieceImpl = (*boltPiece)(nil)
    24  	dataBucketKey           = []byte("data")
    25  )
    26  
    27  func (me *boltPiece) pc() PieceCompletionGetSetter {
    28  	return boltPieceCompletion{me.db}
    29  }
    30  
    31  func (me *boltPiece) pk() metainfo.PieceKey {
    32  	return metainfo.PieceKey{me.ih, me.p.Index()}
    33  }
    34  
    35  func (me *boltPiece) Completion() Completion {
    36  	c, err := me.pc().Get(me.pk())
    37  	switch err {
    38  	case bbolt.ErrDatabaseNotOpen:
    39  		return Completion{}
    40  	case nil:
    41  	default:
    42  		panic(err)
    43  	}
    44  	return c
    45  }
    46  
    47  func (me *boltPiece) MarkComplete() error {
    48  	return me.pc().Set(me.pk(), true)
    49  }
    50  
    51  func (me *boltPiece) MarkNotComplete() error {
    52  	return me.pc().Set(me.pk(), false)
    53  }
    54  
    55  func (me *boltPiece) ReadAt(b []byte, off int64) (n int, err error) {
    56  	err = me.db.View(func(tx *bbolt.Tx) error {
    57  		db := tx.Bucket(dataBucketKey)
    58  		if db == nil {
    59  			return io.EOF
    60  		}
    61  		ci := off / chunkSize
    62  		off %= chunkSize
    63  		for len(b) != 0 {
    64  			ck := me.chunkKey(int(ci))
    65  			_b := db.Get(ck[:])
    66  			// If the chunk is the wrong size, assume it's missing as we can't rely on the data.
    67  			if len(_b) != chunkSize {
    68  				return io.EOF
    69  			}
    70  			n1 := copy(b, _b[off:])
    71  			off = 0
    72  			ci++
    73  			b = b[n1:]
    74  			n += n1
    75  		}
    76  		return nil
    77  	})
    78  	return
    79  }
    80  
    81  func (me *boltPiece) chunkKey(index int) (ret [26]byte) {
    82  	copy(ret[:], me.key[:])
    83  	binary.BigEndian.PutUint16(ret[24:], uint16(index))
    84  	return
    85  }
    86  
    87  func (me *boltPiece) WriteAt(b []byte, off int64) (n int, err error) {
    88  	err = me.db.Update(func(tx *bbolt.Tx) error {
    89  		db, err := tx.CreateBucketIfNotExists(dataBucketKey)
    90  		if err != nil {
    91  			return err
    92  		}
    93  		ci := off / chunkSize
    94  		off %= chunkSize
    95  		for len(b) != 0 {
    96  			_b := make([]byte, chunkSize)
    97  			ck := me.chunkKey(int(ci))
    98  			copy(_b, db.Get(ck[:]))
    99  			n1 := copy(_b[off:], b)
   100  			db.Put(ck[:], _b)
   101  			if n1 > len(b) {
   102  				break
   103  			}
   104  			b = b[n1:]
   105  			off = 0
   106  			ci++
   107  			n += n1
   108  		}
   109  		return nil
   110  	})
   111  	return
   112  }