storage: Fix race condition in mapPieceCompletion

diff --git a/storage/completion.go b/storage/completion.go
index 15fcc249d9a3ebfbf201094866a599c349f779fb..7804351036a238c25b9a3b0ddf06ae312d492dcc 100644 (file)
--- a/storage/completion.go
+++ b/storage/completion.go
@@ -6,7 +6,7 @@ import (
        "github.com/anacrolix/torrent/metainfo"
 )
 
-// Implementations track the completion of pieces.
+// Implementations track the completion of pieces. It must be concurrent-safe.
 type pieceCompletion interface {
        Get(metainfo.PieceKey) (bool, error)
        Set(metainfo.PieceKey, bool) error

diff --git a/storage/completion_piece_map.go b/storage/completion_piece_map.go
index cca1324510ba79d6adff99895ac3e3f1c45a5d63..f4b4e28582d0744b7e46e2c913436c1d323b8187 100644 (file)
--- a/storage/completion_piece_map.go
+++ b/storage/completion_piece_map.go
@@ -1,21 +1,27 @@
 package storage
 
 import (
+       "sync"
+
        "github.com/anacrolix/torrent/metainfo"
 )
 
 type mapPieceCompletion struct {
-       m map[metainfo.PieceKey]struct{}
+       mu sync.Mutex
+       m  map[metainfo.PieceKey]struct{}
 }
 
 func (mapPieceCompletion) Close() error { return nil }
 
 func (me *mapPieceCompletion) Get(pk metainfo.PieceKey) (bool, error) {
+       me.mu.Lock()
        _, ok := me.m[pk]
+       me.mu.Unlock()
        return ok, nil
 }
 
 func (me *mapPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
+       me.mu.Lock()
        if b {
                if me.m == nil {
                        me.m = make(map[metainfo.PieceKey]struct{})
@@ -24,5 +30,6 @@ func (me *mapPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
        } else {
                delete(me.m, pk)
        }
+       me.mu.Unlock()
        return nil
 }