storage/test/bench-resource-pieces.go

   1 package test_storage
   2
   3 import (
   4         "bytes"
   5         "io"
   6         "io/ioutil"
   7         "math/rand"
   8         "sync"
   9         "testing"
  10
  11         "github.com/anacrolix/torrent/metainfo"
  12         "github.com/anacrolix/torrent/storage"
  13         "github.com/bradfitz/iter"
  14         qt "github.com/frankban/quicktest"
  15 )
  16
  17 const (
  18         ChunkSize        = 1 << 14
  19         DefaultPieceSize = 2 << 20
  20         DefaultCapacity  = 0
  21         DefaultNumPieces = 16
  22 )
  23
  24 func BenchmarkPieceMarkComplete(
  25         b *testing.B, ci storage.ClientImpl,
  26         pieceSize int64, numPieces int, capacity int64,
  27 ) {
  28         const check = true
  29         c := qt.New(b)
  30         info := &metainfo.Info{
  31                 Pieces:      make([]byte, numPieces*metainfo.HashSize),
  32                 PieceLength: pieceSize,
  33                 Length:      pieceSize * int64(numPieces),
  34                 Name:        "TorrentName",
  35         }
  36         ti, err := ci.OpenTorrent(info, metainfo.Hash{})
  37         c.Assert(err, qt.IsNil)
  38         defer ti.Close()
  39         rand.Read(info.Pieces)
  40         data := make([]byte, pieceSize)
  41         b.SetBytes(int64(numPieces) * pieceSize)
  42         oneIter := func() {
  43                 for pieceIndex := range iter.N(numPieces) {
  44                         pi := ti.Piece(info.Piece(pieceIndex))
  45                         if check {
  46                                 rand.Read(data)
  47                         }
  48                         var wg sync.WaitGroup
  49                         for off := int64(0); off < int64(len(data)); off += ChunkSize {
  50                                 wg.Add(1)
  51                                 go func(off int64) {
  52                                         defer wg.Done()
  53                                         n, err := pi.WriteAt(data[off:off+ChunkSize], off)
  54                                         if err != nil {
  55                                                 panic(err)
  56                                         }
  57                                         if n != ChunkSize {
  58                                                 panic(n)
  59                                         }
  60                                 }(off)
  61                         }
  62                         wg.Wait()
  63                         if capacity == 0 {
  64                                 pi.MarkNotComplete()
  65                         }
  66                         // This might not apply if users of this benchmark don't cache with the expected capacity.
  67                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
  68                         c.Assert(pi.MarkComplete(), qt.IsNil)
  69                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
  70                         if check {
  71                                 readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
  72                                 c.Assert(err, qt.IsNil)
  73                                 c.Assert(len(readData), qt.Equals, len(data))
  74                                 c.Assert(bytes.Equal(readData, data), qt.IsTrue)
  75                         }
  76                 }
  77         }
  78         // Fill the cache
  79         if capacity > 0 {
  80                 for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
  81                         oneIter()
  82                 }
  83         }
  84         b.ResetTimer()
  85         for range iter.N(b.N) {
  86                 oneIter()
  87         }
  88 }