peer_protocol/decoder_test.go

   1 package peer_protocol
   2
   3 import (
   4         "bufio"
   5         "io"
   6         "sync"
   7         "testing"
   8
   9         "github.com/bradfitz/iter"
  10         "github.com/stretchr/testify/require"
  11 )
  12
  13 func BenchmarkDecodePieces(t *testing.B) {
  14         r, w := io.Pipe()
  15         const pieceLen = 1 << 14
  16         msg := Message{
  17                 Type:  Piece,
  18                 Index: 0,
  19                 Begin: 1,
  20                 Piece: make([]byte, pieceLen),
  21         }
  22         b, err := msg.MarshalBinary()
  23         require.NoError(t, err)
  24         t.SetBytes(int64(len(b)))
  25         defer r.Close()
  26         go func() {
  27                 defer w.Close()
  28                 for {
  29                         n, err := w.Write(b)
  30                         if err == io.ErrClosedPipe {
  31                                 return
  32                         }
  33                         require.NoError(t, err)
  34                         require.Equal(t, len(b), n)
  35                 }
  36         }()
  37         d := Decoder{
  38                 // Emulate what package torrent's client would do.
  39                 R:         bufio.NewReader(r),
  40                 MaxLength: 1 << 18,
  41                 Pool: &sync.Pool{
  42                         New: func() interface{} {
  43                                 b := make([]byte, pieceLen)
  44                                 return &b
  45                         },
  46                 },
  47         }
  48         for range iter.N(t.N) {
  49                 var msg Message
  50                 require.NoError(t, d.Decode(&msg))
  51                 // WWJD
  52                 d.Pool.Put(&msg.Piece)
  53         }
  54 }