storage/bolt-piece-completion.go

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