storage/test/bench-piece-mark-complete.go

   1 package test_storage
   2
   3 import (
   4         "bytes"
   5         "math/rand"
   6         "sync"
   7         "testing"
   8
   9         "github.com/anacrolix/torrent/metainfo"
  10         "github.com/anacrolix/torrent/storage"
  11         "github.com/bradfitz/iter"
  12         qt "github.com/frankban/quicktest"
  13 )
  14
  15 const (
  16         ChunkSize        = 1 << 14
  17         DefaultPieceSize = 2 << 20
  18         DefaultCapacity  = 0
  19         DefaultNumPieces = 16
  20 )
  21
  22 func BenchmarkPieceMarkComplete(
  23         b *testing.B, ci storage.ClientImpl,
  24         pieceSize int64, numPieces int,
  25         // This drives any special handling around capacity that may be configured into the storage
  26         // implementation.
  27         capacity int64,
  28 ) {
  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         tw := storage.Torrent{ti}
  39         defer tw.Close()
  40         rand.Read(info.Pieces)
  41         data := make([]byte, pieceSize)
  42         readData := make([]byte, pieceSize)
  43         b.SetBytes(int64(numPieces) * pieceSize)
  44         oneIter := func() {
  45                 for pieceIndex := range iter.N(numPieces) {
  46                         pi := tw.Piece(info.Piece(pieceIndex))
  47                         rand.Read(data)
  48                         b.StartTimer()
  49                         var wg sync.WaitGroup
  50                         for off := int64(0); off < int64(len(data)); off += ChunkSize {
  51                                 wg.Add(1)
  52                                 go func(off int64) {
  53                                         defer wg.Done()
  54                                         n, err := pi.WriteAt(data[off:off+ChunkSize], off)
  55                                         if err != nil {
  56                                                 panic(err)
  57                                         }
  58                                         if n != ChunkSize {
  59                                                 panic(n)
  60                                         }
  61                                 }(off)
  62                         }
  63                         wg.Wait()
  64                         if capacity == 0 {
  65                                 pi.MarkNotComplete()
  66                         }
  67                         // This might not apply if users of this benchmark don't cache with the expected capacity.
  68                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{Complete: false, Ok: true})
  69                         c.Assert(pi.MarkComplete(), qt.IsNil)
  70                         c.Assert(pi.Completion(), qt.Equals, storage.Completion{true, true})
  71                         n, err := pi.WriteTo(bytes.NewBuffer(readData[:0]))
  72                         b.StopTimer()
  73                         c.Assert(err, qt.IsNil)
  74                         c.Assert(n, qt.Equals, int64(len(data)))
  75                         c.Assert(bytes.Equal(readData[:n], data), qt.IsTrue)
  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 }