storage/sqlite-piece-completion.go

   1 //go:build !nosqlite
   2 // +build !nosqlite
   3
   4 package storage
   5
   6 import (
   7         "errors"
   8         "path/filepath"
   9         "sync"
  10
  11         "github.com/anacrolix/torrent/metainfo"
  12         "zombiezen.com/go/sqlite"
  13         "zombiezen.com/go/sqlite/sqlitex"
  14 )
  15
  16 type sqlitePieceCompletion struct {
  17         mu     sync.Mutex
  18         closed bool
  19         db     *sqlite.Conn
  20 }
  21
  22 var _ PieceCompletion = (*sqlitePieceCompletion)(nil)
  23
  24 func NewSqlitePieceCompletion(dir string) (ret *sqlitePieceCompletion, err error) {
  25         p := filepath.Join(dir, ".torrent.db")
  26         db, err := sqlite.OpenConn(p, 0)
  27         if err != nil {
  28                 return
  29         }
  30         err = sqlitex.ExecScript(db, `create table if not exists piece_completion(infohash, "index", complete, unique(infohash, "index"))`)
  31         if err != nil {
  32                 db.Close()
  33                 return
  34         }
  35         ret = &sqlitePieceCompletion{db: db}
  36         return
  37 }
  38
  39 func (me *sqlitePieceCompletion) Get(pk metainfo.PieceKey) (c Completion, err error) {
  40         me.mu.Lock()
  41         defer me.mu.Unlock()
  42         err = sqlitex.Exec(
  43                 me.db, `select complete from piece_completion where infohash=? and "index"=?`,
  44                 func(stmt *sqlite.Stmt) error {
  45                         c.Complete = stmt.ColumnInt(0) != 0
  46                         c.Ok = true
  47                         return nil
  48                 },
  49                 pk.InfoHash.HexString(), pk.Index)
  50         return
  51 }
  52
  53 func (me *sqlitePieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
  54         me.mu.Lock()
  55         defer me.mu.Unlock()
  56         if me.closed {
  57                 return errors.New("closed")
  58         }
  59         return sqlitex.Exec(
  60                 me.db,
  61                 `insert or replace into piece_completion(infohash, "index", complete) values(?, ?, ?)`,
  62                 nil,
  63                 pk.InfoHash.HexString(), pk.Index, b)
  64 }
  65
  66 func (me *sqlitePieceCompletion) Close() (err error) {
  67         me.mu.Lock()
  68         defer me.mu.Unlock()
  69         if me.closed {
  70                 return
  71         }
  72         err = me.db.Close()
  73         me.db = nil
  74         me.closed = true
  75         return
  76 }