From cee200a5a217d930ef955d352709f0edb81b7007 Mon Sep 17 00:00:00 2001 From: Matt Joiner Date: Mon, 20 Jun 2016 17:51:05 +1000 Subject: [PATCH] Add piece completion storage Toward fixing https://github.com/anacrolix/torrent/issues/50. --- client.go | 3 ++ metainfo/piece.go | 13 ++++++++ storage/completion.go | 12 ++++++++ storage/completion_piece_map.go | 27 +++++++++++++++++ storage/db.go | 54 +++++++++++++++++++++++++++++++++ storage/file.go | 22 +++++++++----- 6 files changed, 123 insertions(+), 8 deletions(-) create mode 100644 storage/completion.go create mode 100644 storage/completion_piece_map.go create mode 100644 storage/db.go diff --git a/client.go b/client.go index 8c388280..bb7c0976 100644 --- 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 7d952126..a68f1bce 100644 --- 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 index 00000000..99dd99a9 --- /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 index 00000000..8b904fb9 --- /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 index 00000000..5d54c882 --- /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 f7815753..a3a26bed 100644 --- 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 } -- 2.44.0