peer_protocol/decoder_test.go

   1 package peer_protocol
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "io"
   7         "sync"
   8         "testing"
   9
  10         qt "github.com/frankban/quicktest"
  11         "github.com/stretchr/testify/assert"
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 func BenchmarkDecodePieces(t *testing.B) {
  16         const pieceLen = 1 << 14
  17         inputMsg := Message{
  18                 Type:  Piece,
  19                 Index: 0,
  20                 Begin: 1,
  21                 Piece: make([]byte, pieceLen),
  22         }
  23         b := inputMsg.MustMarshalBinary()
  24         t.SetBytes(int64(len(b)))
  25         var r bytes.Reader
  26         // Try to somewhat emulate what torrent.Client would do. But the goal is to get decoding as fast
  27         // as possible and let consumers apply their own adjustments.
  28         d := Decoder{
  29                 R:         bufio.NewReaderSize(&r, 1<<10),
  30                 MaxLength: 1 << 18,
  31                 Pool: &sync.Pool{
  32                         New: func() interface{} {
  33                                 b := make([]byte, pieceLen)
  34                                 return &b
  35                         },
  36                 },
  37         }
  38         c := qt.New(t)
  39         t.ReportAllocs()
  40         t.ResetTimer()
  41         for i := 0; i < t.N; i += 1 {
  42                 r.Reset(b)
  43                 var msg Message
  44                 err := d.Decode(&msg)
  45                 if err != nil {
  46                         t.Fatal(err)
  47                 }
  48                 // This is very expensive, and should be discovered in tests rather than a benchmark.
  49                 if false {
  50                         c.Assert(msg, qt.DeepEquals, inputMsg)
  51                 }
  52                 // WWJD
  53                 d.Pool.Put(&msg.Piece)
  54         }
  55 }
  56
  57 func TestDecodeShortPieceEOF(t *testing.T) {
  58         r, w := io.Pipe()
  59         go func() {
  60                 w.Write(Message{Type: Piece, Piece: make([]byte, 1)}.MustMarshalBinary())
  61                 w.Close()
  62         }()
  63         d := Decoder{
  64                 R:         bufio.NewReader(r),
  65                 MaxLength: 1 << 15,
  66                 Pool: &sync.Pool{New: func() interface{} {
  67                         b := make([]byte, 2)
  68                         return &b
  69                 }},
  70         }
  71         var m Message
  72         require.NoError(t, d.Decode(&m))
  73         assert.Len(t, m.Piece, 1)
  74         assert.ErrorIs(t, d.Decode(&m), io.EOF)
  75 }
  76
  77 func TestDecodeOverlongPiece(t *testing.T) {
  78         r, w := io.Pipe()
  79         go func() {
  80                 w.Write(Message{Type: Piece, Piece: make([]byte, 3)}.MustMarshalBinary())
  81                 w.Close()
  82         }()
  83         d := Decoder{
  84                 R:         bufio.NewReader(r),
  85                 MaxLength: 1 << 15,
  86                 Pool: &sync.Pool{New: func() interface{} {
  87                         b := make([]byte, 2)
  88                         return &b
  89                 }},
  90         }
  91         var m Message
  92         require.Error(t, d.Decode(&m))
  93 }