storage/bolt-piece.go

   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 }