storage/bolt-piece-completion.go

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