github.com/anacrolix/torrent@v1.61.0/storage/bolt-piece-completion.go (about)

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