From 0021c2a70c2a3ebb5c33d9efcd0048be8d3695d9 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 28 Apr 2021 14:46:47 +1000
Subject: [PATCH] Use BenchmarkMarkComplete for non-sqlite storages too

---
 storage/mark-complete_test.go         | 30 ++++++++++++++++++
 storage/sqlite/sqlite-storage_test.go |  6 ++--
 storage/test/bench-resource-pieces.go | 44 +++++++++++++++++----------
 3 files changed, 61 insertions(+), 19 deletions(-)
 create mode 100644 storage/mark-complete_test.go

diff --git a/storage/mark-complete_test.go b/storage/mark-complete_test.go
new file mode 100644
index 00000000..3f76f699
--- /dev/null
+++ b/storage/mark-complete_test.go
@@ -0,0 +1,30 @@
+package storage_test
+
+import (
+	"testing"
+
+	"github.com/anacrolix/torrent/storage"
+	test_storage "github.com/anacrolix/torrent/storage/test"
+)
+
+func BenchmarkMarkComplete(b *testing.B) {
+	bench := func(b *testing.B, ci storage.ClientImpl) {
+		test_storage.BenchmarkPieceMarkComplete(
+			b, ci, test_storage.DefaultPieceSize, test_storage.DefaultNumPieces, test_storage.DefaultCapacity)
+	}
+	b.Run("File", func(b *testing.B) {
+		ci := storage.NewFile(b.TempDir())
+		b.Cleanup(func() { ci.Close() })
+		bench(b, ci)
+	})
+	b.Run("Mmap", func(b *testing.B) {
+		ci := storage.NewMMap(b.TempDir())
+		b.Cleanup(func() { ci.Close() })
+		bench(b, ci)
+	})
+	b.Run("BoltDb", func(b *testing.B) {
+		ci := storage.NewBoltDB(b.TempDir())
+		b.Cleanup(func() { ci.Close() })
+		bench(b, ci)
+	})
+}
diff --git a/storage/sqlite/sqlite-storage_test.go b/storage/sqlite/sqlite-storage_test.go
index 3f599d51..a903e3a9 100644
--- a/storage/sqlite/sqlite-storage_test.go
+++ b/storage/sqlite/sqlite-storage_test.go
@@ -67,8 +67,8 @@ func TestSimultaneousIncrementalBlob(t *testing.T) {
 }
 
 func BenchmarkMarkComplete(b *testing.B) {
-	const pieceSize = 2 << 20
-	const capacity = 0
+	const pieceSize = test_storage.DefaultPieceSize
+	const capacity = test_storage.DefaultCapacity
 	c := qt.New(b)
 	for _, memory := range []bool{false, true} {
 		b.Run(fmt.Sprintf("Memory=%v", memory), func(b *testing.B) {
@@ -90,7 +90,7 @@ func BenchmarkMarkComplete(b *testing.B) {
 					})
 					c.Assert(err, qt.IsNil)
 					defer ci.Close()
-					test_storage.BenchmarkPieceMarkComplete(b, ci, pieceSize, 16, capacity)
+					test_storage.BenchmarkPieceMarkComplete(b, ci, pieceSize, test_storage.DefaultNumPieces, capacity)
 				})
 			}
 		})
diff --git a/storage/test/bench-resource-pieces.go b/storage/test/bench-resource-pieces.go
index 95646484..4a79c313 100644
--- a/storage/test/bench-resource-pieces.go
+++ b/storage/test/bench-resource-pieces.go
@@ -14,37 +14,47 @@ import (
 	qt "github.com/frankban/quicktest"
 )
 
-const chunkSize = 1 << 14
+const (
+	ChunkSize        = 1 << 14
+	DefaultPieceSize = 2 << 20
+	DefaultCapacity  = 0
+	DefaultNumPieces = 16
+)
 
-func BenchmarkPieceMarkComplete(b *testing.B, ci storage.ClientImpl, pieceSize int64, numPieces int, capacity int64) {
+func BenchmarkPieceMarkComplete(
+	b *testing.B, ci storage.ClientImpl,
+	pieceSize int64, numPieces int, capacity int64,
+) {
+	const check = true
 	c := qt.New(b)
-	ti, err := ci.OpenTorrent(&metainfo.Info{
-		Pieces:      make([]byte, metainfo.HashSize*numPieces),
-		PieceLength: pieceSize,
-	}, metainfo.Hash{})
-	c.Assert(err, qt.IsNil)
-	defer ti.Close()
 	info := &metainfo.Info{
 		Pieces:      make([]byte, numPieces*metainfo.HashSize),
 		PieceLength: pieceSize,
 		Length:      pieceSize * int64(numPieces),
+		Name:        "TorrentName",
 	}
+	ti, err := ci.OpenTorrent(info, metainfo.Hash{})
+	c.Assert(err, qt.IsNil)
+	defer ti.Close()
 	rand.Read(info.Pieces)
 	data := make([]byte, pieceSize)
+	b.SetBytes(int64(numPieces) * pieceSize)
 	oneIter := func() {
 		for pieceIndex := range iter.N(numPieces) {
 			pi := ti.Piece(info.Piece(pieceIndex))
-			rand.Read(data)
+			if check {
+				rand.Read(data)
+			}
 			var wg sync.WaitGroup
-			for off := int64(0); off < int64(len(data)); off += chunkSize {
+			for off := int64(0); off < int64(len(data)); off += ChunkSize {
 				wg.Add(1)
 				go func(off int64) {
 					defer wg.Done()
-					n, err := pi.WriteAt(data[off:off+chunkSize], off)
+					n, err := pi.WriteAt(data[off:off+ChunkSize], off)
 					if err != nil {
 						panic(err)
 					}
-					if n != chunkSize {
+					if n != ChunkSize {
 						panic(n)
 					}
 				}(off)
@@ -57,10 +67,12 @@ func BenchmarkPieceMarkComplete(b *testing.B, ci storage.ClientImpl, pieceSize i
 			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})
-			readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
-			c.Assert(err, qt.IsNil)
-			c.Assert(len(readData), qt.Equals, len(data))
-			c.Assert(bytes.Equal(readData, data), qt.IsTrue)
+			if check {
+				readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
+				c.Assert(err, qt.IsNil)
+				c.Assert(len(readData), qt.Equals, len(data))
+				c.Assert(bytes.Equal(readData, data), qt.IsTrue)
+			}
 		}
 	}
 	// Fill the cache
-- 
2.44.0