storage/bolt-piece.go

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