Add piece completion storage

Toward fixing https://github.com/anacrolix/torrent/issues/50.

diff --git a/client.go b/client.go
index 8c388280576f9c983591e093d81e14d2099576a4..bb7c0976c960dc602f395e5de08fa9156b3ff8f1 100644 (file)
--- a/client.go
+++ b/client.go
@@ -1705,6 +1705,9 @@ func (cl *Client) reapPieceTouches(t *Torrent, piece int) (ret []*connection) {
 }
 
 func (cl *Client) pieceHashed(t *Torrent, piece int, correct bool) {
+       if t.closed.IsSet() {
+               return
+       }
        p := &t.pieces[piece]
        if p.EverHashed {
                // Don't score the first time a piece is hashed, it could be an

diff --git a/metainfo/piece.go b/metainfo/piece.go
index 7d95212671c675967bbece0d7fa24bbd8e935973..a68f1bce7f538c196b24dbde6bedd64005dcd1a7 100644 (file)
--- a/metainfo/piece.go
+++ b/metainfo/piece.go
@@ -22,3 +22,16 @@ func (p Piece) Hash() (ret Hash) {
        missinggo.CopyExact(&ret, p.Info.Pieces[p.i*20:(p.i+1)*20])
        return
 }
+
+func (p Piece) Index() int {
+       return p.i
+}
+
+func (p Piece) Key() PieceKey {
+       return PieceKey{p.Info.Hash(), p.i}
+}
+
+type PieceKey struct {
+       Hash  Hash
+       Index int
+}

diff --git a/storage/completion.go b/storage/completion.go
new file mode 100644 (file)
index 0000000..99dd99a
--- /dev/null
+++ b/storage/completion.go
@@ -0,0 +1,12 @@
+package storage
+
+import "log"
+
+func pieceCompletionForDir(dir string) (ret pieceCompletion) {
+       ret, err := newDBPieceCompletion(dir)
+       if err != nil {
+               log.Printf("couldn't open piece completion db in %q: %s", dir, err)
+               ret = new(mapPieceCompletion)
+       }
+       return
+}

diff --git a/storage/completion_piece_map.go b/storage/completion_piece_map.go
new file mode 100644 (file)
index 0000000..8b904fb
--- /dev/null
+++ b/storage/completion_piece_map.go
@@ -0,0 +1,27 @@
+package storage
+
+import (
+       "github.com/anacrolix/torrent/metainfo"
+)
+
+type mapPieceCompletion struct {
+       m map[metainfo.PieceKey]struct{}
+}
+
+func (mapPieceCompletion) Close() {}
+
+func (me *mapPieceCompletion) Get(p metainfo.Piece) bool {
+       _, ok := me.m[p.Key()]
+       return ok
+}
+
+func (me *mapPieceCompletion) Set(p metainfo.Piece, b bool) {
+       if b {
+               if me.m == nil {
+                       me.m = make(map[metainfo.PieceKey]struct{})
+               }
+               me.m[p.Key()] = struct{}{}
+       } else {
+               delete(me.m, p.Key())
+       }
+}

diff --git a/storage/db.go b/storage/db.go
new file mode 100644 (file)
index 0000000..5d54c88
--- /dev/null
+++ b/storage/db.go
@@ -0,0 +1,54 @@
+package storage
+
+import (
+       "database/sql"
+       "log"
+       "path/filepath"
+
+       "github.com/anacrolix/torrent/metainfo"
+)
+
+type dbPieceCompletion struct {
+       db *sql.DB
+}
+
+func newDBPieceCompletion(dir string) (ret *dbPieceCompletion, err error) {
+       p := filepath.Join(dir, ".torrent.db")
+       log.Println(p)
+       db, err := sql.Open("sqlite3", p)
+       if err != nil {
+               return
+       }
+       _, err = db.Exec(`create table if not exists completed(infohash, "index", unique(infohash, "index") on conflict ignore)`)
+       if err != nil {
+               db.Close()
+               return
+       }
+       ret = &dbPieceCompletion{db}
+       return
+}
+
+func (me *dbPieceCompletion) Get(p metainfo.Piece) (ret bool) {
+       row := me.db.QueryRow(`select exists(select * from completed where infohash=? and "index"=?)`, p.Info.Hash().HexString(), p.Index())
+       err := row.Scan(&ret)
+       if err != nil {
+               panic(err)
+       }
+       return
+}
+
+func (me *dbPieceCompletion) Set(p metainfo.Piece, b bool) {
+       var err error
+       if b {
+               _, err = me.db.Exec(`insert into completed (infohash, "index") values (?, ?)`, p.Info.Hash().HexString(), p.Index())
+       } else {
+               _, err = me.db.Exec(`delete from completed where infohash=? and "index"=?`, p.Info.Hash().HexString(), p.Index())
+       }
+       if err != nil {
+               panic(err)
+       }
+}
+
+func (me *dbPieceCompletion) Close() {
+       me.db.Close()
+}

diff --git a/storage/file.go b/storage/file.go
index f781575382e1a200147b17907e46069ad5f48edf..a3a26bede2954abb68c417c19347058650b07629 100644 (file)
--- a/storage/file.go
+++ b/storage/file.go
@@ -6,20 +6,28 @@ import (
        "path/filepath"
 
        "github.com/anacrolix/missinggo"
+       _ "github.com/mattn/go-sqlite3"
 
        "github.com/anacrolix/torrent/metainfo"
 )
 
+type pieceCompletion interface {
+       Get(metainfo.Piece) bool
+       Set(metainfo.Piece, bool)
+       Close()
+}
+
 // File-based storage for torrents, that isn't yet bound to a particular
 // torrent.
 type fileStorage struct {
-       baseDir   string
-       completed map[[20]byte]bool
+       baseDir    string
+       completion pieceCompletion
 }
 
 func NewFile(baseDir string) Client {
        return &fileStorage{
-               baseDir: baseDir,
+               baseDir:    baseDir,
+               completion: pieceCompletionForDir(baseDir),
        }
 }
 
@@ -48,6 +56,7 @@ func (fs *fileStorage) Piece(p metainfo.Piece) Piece {
 }
 
 func (fs *fileStorage) Close() error {
+       fs.completion.Close()
        return nil
 }
 
@@ -59,14 +68,11 @@ type fileStoragePiece struct {
 }
 
 func (fs *fileStoragePiece) GetIsComplete() bool {
-       return fs.completed[fs.p.Hash()]
+       return fs.completion.Get(fs.p)
 }
 
 func (fs *fileStoragePiece) MarkComplete() error {
-       if fs.completed == nil {
-               fs.completed = make(map[[20]byte]bool)
-       }
-       fs.completed[fs.p.Hash()] = true
+       fs.completion.Set(fs.p, true)
        return nil
 }