storage/boltPieceCompletion.go

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