Fix UseSources panicking when sqlite storage is closed

[btrtrc.git] / storage / test / bench-piece-mark-complete.go

diff --git a/storage/test/bench-piece-mark-complete.go b/storage/test/bench-piece-mark-complete.go
index 1656b8acbdc22d941c335355aa8416704f62f75d..db0049f22674b98ff499bd12745799133d01869a 100644 (file)
--- a/storage/test/bench-piece-mark-complete.go
+++ b/storage/test/bench-piece-mark-complete.go
@@ -2,22 +2,19 @@ package test_storage
 
 import (
        "bytes"
-       "io"
-       "io/ioutil"
        "math/rand"
        "sync"
        "testing"
 
+       qt "github.com/frankban/quicktest"
+
        "github.com/anacrolix/torrent/metainfo"
        "github.com/anacrolix/torrent/storage"
-       "github.com/bradfitz/iter"
-       qt "github.com/frankban/quicktest"
 )
 
 const (
        ChunkSize        = 1 << 14
        DefaultPieceSize = 2 << 20
-       DefaultCapacity  = 0
        DefaultNumPieces = 16
 )
 
@@ -28,7 +25,6 @@ func BenchmarkPieceMarkComplete(
        // implementation.
        capacity int64,
 ) {
-       const check = true
        c := qt.New(b)
        info := &metainfo.Info{
                Pieces:      make([]byte, numPieces*metainfo.HashSize),
@@ -38,16 +34,17 @@ func BenchmarkPieceMarkComplete(
        }
        ti, err := ci.OpenTorrent(info, metainfo.Hash{})
        c.Assert(err, qt.IsNil)
-       defer ti.Close()
+       tw := storage.Torrent{ti}
+       defer tw.Close()
        rand.Read(info.Pieces)
        data := make([]byte, pieceSize)
+       readData := make([]byte, pieceSize)
        b.SetBytes(int64(numPieces) * pieceSize)
        oneIter := func() {
-               for pieceIndex := range iter.N(numPieces) {
-                       pi := ti.Piece(info.Piece(pieceIndex))
-                       if check {
-                               rand.Read(data)
-                       }
+               for pieceIndex := 0; pieceIndex < numPieces; pieceIndex += 1 {
+                       pi := tw.Piece(info.Piece(pieceIndex))
+                       rand.Read(data)
+                       b.StartTimer()
                        var wg sync.WaitGroup
                        for off := int64(0); off < int64(len(data)); off += ChunkSize {
                                wg.Add(1)
@@ -63,31 +60,29 @@ func BenchmarkPieceMarkComplete(
                                }(off)
                        }
                        wg.Wait()
-                       b.StopTimer()
                        if capacity == 0 {
                                pi.MarkNotComplete()
                        }
-                       b.StartTimer()
                        // This might not apply if users of this benchmark don't cache with the expected capacity.
                        c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
                        c.Assert(pi.MarkComplete(), qt.IsNil)
-                       c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
-                       if check {
-                               readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
-                               c.Check(err, qt.IsNil)
-                               c.Assert(len(readData), qt.Equals, len(data))
-                               c.Assert(bytes.Equal(readData, data), qt.IsTrue)
-                       }
+                       c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: true, Ok: true})
+                       n, err := pi.WriteTo(bytes.NewBuffer(readData[:0]))
+                       b.StopTimer()
+                       c.Assert(err, qt.IsNil)
+                       c.Assert(n, qt.Equals, int64(len(data)))
+                       c.Assert(bytes.Equal(readData[:n], data), qt.IsTrue)
                }
        }
        // Fill the cache
        if capacity > 0 {
-               for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
+               iterN := int((capacity + info.TotalLength() - 1) / info.TotalLength())
+               for i := 0; i < iterN; i += 1 {
                        oneIter()
                }
        }
        b.ResetTimer()
-       for range iter.N(b.N) {
+       for i := 0; i < b.N; i += 1 {
                oneIter()
        }
 }