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         qt "github.com/frankban/quicktest"
  10
  11         "github.com/anacrolix/torrent/metainfo"
  12         "github.com/anacrolix/torrent/storage"
  13 )
  14
  15 const (
  16         ChunkSize        = 1 << 14
  17         DefaultPieceSize = 2 << 20
  18         DefaultNumPieces = 16
  19 )
  20
  21 func BenchmarkPieceMarkComplete(
  22         b *testing.B, ci storage.ClientImpl,
  23         pieceSize int64, numPieces int,
  24         // This drives any special handling around capacity that may be configured into the storage
  25         // implementation.
  26         capacity int64,
  27 ) {
  28         c := qt.New(b)
  29         info := &metainfo.Info{
  30                 Pieces:      make([]byte, numPieces*metainfo.HashSize),
  31                 PieceLength: pieceSize,
  32                 Length:      pieceSize * int64(numPieces),
  33                 Name:        "TorrentName",
  34         }
  35         ti, err := ci.OpenTorrent(info, metainfo.Hash{})
  36         c.Assert(err, qt.IsNil)
  37         tw := storage.Torrent{ti}
  38         defer tw.Close()
  39         rand.Read(info.Pieces)
  40         data := make([]byte, pieceSize)
  41         readData := make([]byte, pieceSize)
  42         b.SetBytes(int64(numPieces) * pieceSize)
  43         oneIter := func() {
  44                 for pieceIndex := 0; pieceIndex < numPieces; pieceIndex += 1 {
  45                         pi := tw.Piece(info.Piece(pieceIndex))
  46                         rand.Read(data)
  47                         b.StartTimer()
  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{Complete: true, Ok: true})
  70                         n, err := pi.WriteTo(bytes.NewBuffer(readData[:0]))
  71                         b.StopTimer()
  72                         c.Assert(err, qt.IsNil)
  73                         c.Assert(n, qt.Equals, int64(len(data)))
  74                         c.Assert(bytes.Equal(readData[:n], data), qt.IsTrue)
  75                 }
  76         }
  77         // Fill the cache
  78         if capacity > 0 {
  79                 iterN := int((capacity + info.TotalLength() - 1) / info.TotalLength())
  80                 for i := 0; i < iterN; i += 1 {
  81                         oneIter()
  82                 }
  83         }
  84         b.ResetTimer()
  85         for i := 0; i < b.N; i += 1 {
  86                 oneIter()
  87         }
  88 }