storage/sqlite-piece-completion.go

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