storage/bolt-piece.go

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