Add piece completion storage

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

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

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
+}

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
+}

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())
+	}
+}

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()
+}

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
 }