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 completionBucketKey = []byte("completion")
  23
  24 type boltPieceCompletion struct {
  25         db *bbolt.DB
  26 }
  27
  28 var _ PieceCompletion = (*boltPieceCompletion)(nil)
  29
  30 func NewBoltPieceCompletion(dir string) (ret PieceCompletion, err error) {
  31         os.MkdirAll(dir, 0o750)
  32         p := filepath.Join(dir, ".torrent.bolt.db")
  33         db, err := bbolt.Open(p, 0o660, &bbolt.Options{
  34                 Timeout: time.Second,
  35         })
  36         if err != nil {
  37                 return
  38         }
  39         db.NoSync = true
  40         ret = &boltPieceCompletion{db}
  41         return
  42 }
  43
  44 func (me boltPieceCompletion) Get(pk metainfo.PieceKey) (cn Completion, err error) {
  45         err = me.db.View(func(tx *bbolt.Tx) error {
  46                 cb := tx.Bucket(completionBucketKey)
  47                 if cb == nil {
  48                         return nil
  49                 }
  50                 ih := cb.Bucket(pk.InfoHash[:])
  51                 if ih == nil {
  52                         return nil
  53                 }
  54                 var key [4]byte
  55                 binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
  56                 cn.Ok = true
  57                 switch string(ih.Get(key[:])) {
  58                 case boltDbCompleteValue:
  59                         cn.Complete = true
  60                 case boltDbIncompleteValue:
  61                         cn.Complete = false
  62                 default:
  63                         cn.Ok = false
  64                 }
  65                 return nil
  66         })
  67         return
  68 }
  69
  70 func (me boltPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
  71         if c, err := me.Get(pk); err == nil && c.Ok && c.Complete == b {
  72                 return nil
  73         }
  74         return me.db.Update(func(tx *bbolt.Tx) error {
  75                 c, err := tx.CreateBucketIfNotExists(completionBucketKey)
  76                 if err != nil {
  77                         return err
  78                 }
  79                 ih, err := c.CreateBucketIfNotExists(pk.InfoHash[:])
  80                 if err != nil {
  81                         return err
  82                 }
  83                 var key [4]byte
  84                 binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
  85                 return ih.Put(key[:], []byte(func() string {
  86                         if b {
  87                                 return boltDbCompleteValue
  88                         } else {
  89                                 return boltDbIncompleteValue
  90                         }
  91                 }()))
  92         })
  93 }
  94
  95 func (me *boltPieceCompletion) Close() error {
  96         return me.db.Close()
  97 }