storage/bolt-piece-completion.go

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