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 chunkSize = 1 << 14
  18
  19 func BenchmarkPieceMarkComplete(b *testing.B, ci storage.ClientImpl, pieceSize int64, numPieces int, capacity int64) {
  20         c := qt.New(b)
  21         ti, err := ci.OpenTorrent(&metainfo.Info{
  22                 Pieces:      make([]byte, metainfo.HashSize*numPieces),
  23                 PieceLength: pieceSize,
  24         }, metainfo.Hash{})
  25         c.Assert(err, qt.IsNil)
  26         defer ti.Close()
  27         info := &metainfo.Info{
  28                 Pieces:      make([]byte, numPieces*metainfo.HashSize),
  29                 PieceLength: pieceSize,
  30                 Length:      pieceSize * int64(numPieces),
  31         }
  32         rand.Read(info.Pieces)
  33         data := make([]byte, pieceSize)
  34         oneIter := func() {
  35                 for pieceIndex := range iter.N(numPieces) {
  36                         pi := ti.Piece(info.Piece(pieceIndex))
  37                         rand.Read(data)
  38                         var wg sync.WaitGroup
  39                         for off := int64(0); off < int64(len(data)); off += chunkSize {
  40                                 wg.Add(1)
  41                                 go func(off int64) {
  42                                         defer wg.Done()
  43                                         n, err := pi.WriteAt(data[off:off+chunkSize], off)
  44                                         if err != nil {
  45                                                 panic(err)
  46                                         }
  47                                         if n != chunkSize {
  48                                                 panic(n)
  49                                         }
  50                                 }(off)
  51                         }
  52                         wg.Wait()
  53                         if capacity == 0 {
  54                                 pi.MarkNotComplete()
  55                         }
  56                         // This might not apply if users of this benchmark don't cache with the expected capacity.
  57                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
  58                         c.Assert(pi.MarkComplete(), qt.IsNil)
  59                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
  60                         readData, err := ioutil.ReadAll(io.NewSectionReader(pi, 0, int64(len(data))))
  61                         c.Assert(err, qt.IsNil)
  62                         c.Assert(len(readData), qt.Equals, len(data))
  63                         c.Assert(bytes.Equal(readData, data), qt.IsTrue)
  64                 }
  65         }
  66         // Fill the cache
  67         if capacity > 0 {
  68                 for range iter.N(int((capacity + info.TotalLength() - 1) / info.TotalLength())) {
  69                         oneIter()
  70                 }
  71         }
  72         b.ResetTimer()
  73         for range iter.N(b.N) {
  74                 oneIter()
  75         }
  76 }