storage/bolt_piece.go

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