storage/boltPieceCompletion.go

   1 package storage
   2
   3 import (
   4         "encoding/binary"
   5         "os"
   6         "path/filepath"
   7         "time"
   8
   9         "github.com/boltdb/bolt"
  10
  11         "github.com/anacrolix/torrent/metainfo"
  12 )
  13
  14 var (
  15         value = []byte{}
  16 )
  17
  18 type boltPieceCompletion struct {
  19         db *bolt.DB
  20 }
  21
  22 func NewBoltPieceCompletion(dir string) (ret PieceCompletion, err error) {
  23         os.MkdirAll(dir, 0770)
  24         p := filepath.Join(dir, ".torrent.bolt.db")
  25         db, err := bolt.Open(p, 0660, &bolt.Options{
  26                 Timeout: time.Second,
  27         })
  28         if err != nil {
  29                 return
  30         }
  31         ret = &boltPieceCompletion{db}
  32         return
  33 }
  34
  35 func (me *boltPieceCompletion) Get(pk metainfo.PieceKey) (ret bool, err error) {
  36         err = me.db.View(func(tx *bolt.Tx) error {
  37                 c := tx.Bucket(completed)
  38                 if c == nil {
  39                         return nil
  40                 }
  41                 ih := c.Bucket(pk.InfoHash[:])
  42                 if ih == nil {
  43                         return nil
  44                 }
  45                 var key [4]byte
  46                 binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
  47                 ret = ih.Get(key[:]) != nil
  48                 return nil
  49         })
  50         return
  51 }
  52
  53 func (me *boltPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
  54         return me.db.Update(func(tx *bolt.Tx) error {
  55                 c, err := tx.CreateBucketIfNotExists(completed)
  56                 if err != nil {
  57                         return err
  58                 }
  59                 ih, err := c.CreateBucketIfNotExists(pk.InfoHash[:])
  60                 if err != nil {
  61                         return err
  62                 }
  63                 var key [4]byte
  64                 binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
  65                 if b {
  66                         return ih.Put(key[:], value)
  67                 } else {
  68                         return ih.Delete(key[:])
  69                 }
  70         })
  71 }
  72
  73 func (me *boltPieceCompletion) Close() error {
  74         return me.db.Close()
  75 }