storage/sqlite-piece-completion.go

   1 //go:build cgo
   2 // +build cgo
   3
   4 package storage
   5
   6 import (
   7         "path/filepath"
   8         "sync"
   9
  10         "crawshaw.io/sqlite"
  11         "crawshaw.io/sqlite/sqlitex"
  12
  13         "github.com/anacrolix/torrent/metainfo"
  14 )
  15
  16 type sqlitePieceCompletion struct {
  17         mu sync.Mutex
  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         return sqlitex.Exec(
  56                 me.db,
  57                 `insert or replace into piece_completion(infohash, "index", complete) values(?, ?, ?)`,
  58                 nil,
  59                 pk.InfoHash.HexString(), pk.Index, b)
  60 }
  61
  62 func (me *sqlitePieceCompletion) Close() (err error) {
  63         me.mu.Lock()
  64         defer me.mu.Unlock()
  65         if me.db != nil {
  66                 err = me.db.Close()
  67                 me.db = nil
  68         }
  69         return
  70 }