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