storage/bolt_piece.go

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