storage/test/bench-piece-mark-complete.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,
  27         // This drives any special handling around capacity that may be configured into the storage
  28         // implementation.
  29         capacity int64,
  30 ) {
  31         const check = true
  32         c := qt.New(b)
  33         info := &metainfo.Info{
  34                 Pieces:      make([]byte, numPieces*metainfo.HashSize),
  35                 PieceLength: pieceSize,
  36                 Length:      pieceSize * int64(numPieces),
  37                 Name:        "TorrentName",
  38         }
  39         ti, err := ci.OpenTorrent(info, metainfo.Hash{})
  40         c.Assert(err, qt.IsNil)
  41         defer ti.Close()
  42         rand.Read(info.Pieces)
  43         data := make([]byte, pieceSize)
  44         b.SetBytes(int64(numPieces) * pieceSize)
  45         oneIter := func() {
  46                 for pieceIndex := range iter.N(numPieces) {
  47                         pi := ti.Piece(info.Piece(pieceIndex))
  48                         if check {
  49                                 rand.Read(data)
  50                         }
  51                         var wg sync.WaitGroup
  52                         for off := int64(0); off < int64(len(data)); off += ChunkSize {
  53                                 wg.Add(1)
  54                                 go func(off int64) {
  55                                         defer wg.Done()
  56                                         n, err := pi.WriteAt(data[off:off+ChunkSize], off)
  57                                         if err != nil {
  58                                                 panic(err)
  59                                         }
  60                                         if n != ChunkSize {
  61                                                 panic(n)
  62                                         }
  63                                 }(off)
  64                         }
  65                         wg.Wait()
  66                         b.StopTimer()
  67                         if capacity == 0 {
  68                                 pi.MarkNotComplete()
  69                         }
  70                         b.StartTimer()
  71                         // This might not apply if users of this benchmark don't cache with the expected capacity.
  72                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
  73                         c.Assert(pi.MarkComplete(), qt.IsNil)
  74                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
  75                         if check {
  76                                 readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
  77                                 c.Check(err, qt.IsNil)
  78                                 c.Assert(len(readData), qt.Equals, len(data))
  79                                 c.Assert(bytes.Equal(readData, data), qt.IsTrue)
  80                         }
  81                 }
  82         }
  83         // Fill the cache
  84         if capacity > 0 {
  85                 for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
  86                         oneIter()
  87                 }
  88         }
  89         b.ResetTimer()
  90         for range iter.N(b.N) {
  91                 oneIter()
  92         }
  93 }