github.com/grafana/pyroscope@v1.18.0/pkg/iter/batch_async_test.go (about)

     1  package iter
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  )
     8  
     9  func Test_AsyncBatchIter(t *testing.T) {
    10  	type testCase struct {
    11  		description string
    12  		seqSize     int
    13  		bufSize     int
    14  	}
    15  	testCases := []testCase{
    16  		{
    17  			description: "empty iterator",
    18  			seqSize:     0,
    19  			bufSize:     1,
    20  		},
    21  		{
    22  			description: "empty iterator, zero buffer",
    23  			seqSize:     0,
    24  			bufSize:     0,
    25  		},
    26  		{
    27  			description: "zero buffer",
    28  			seqSize:     10,
    29  			bufSize:     0,
    30  		},
    31  		{
    32  			description: "iterator < buffer",
    33  			seqSize:     5,
    34  			bufSize:     10,
    35  		},
    36  		{
    37  			description: "iterator == buffer",
    38  			seqSize:     10,
    39  			bufSize:     10,
    40  		},
    41  		{
    42  			description: "iterator > buffer",
    43  			seqSize:     25,
    44  			bufSize:     10,
    45  		},
    46  	}
    47  	for _, tc := range testCases {
    48  		tc := tc
    49  		t.Run(tc.description, func(t *testing.T) {
    50  			x := NewAsyncBatchIterator[int, int](
    51  				newSeqIterator(tc.seqSize),
    52  				tc.bufSize,
    53  				func(i int) int { return i },
    54  				func([]int) {},
    55  			)
    56  			var p, c int
    57  			for x.Next() {
    58  				i := x.At()
    59  				require.Equal(t, 1, i-p)
    60  				p = i
    61  				c++
    62  			}
    63  			require.Equal(t, tc.seqSize, c)
    64  			require.NoError(t, x.Err())
    65  			require.NoError(t, x.Close())
    66  		})
    67  	}
    68  }
    69  
    70  type seqIterator struct{ n, c int }
    71  
    72  func newSeqIterator(n int) *seqIterator {
    73  	return &seqIterator{n: n}
    74  }
    75  
    76  func (x *seqIterator) Next() bool {
    77  	if x.c < x.n {
    78  		x.c++
    79  		return true
    80  	}
    81  	return false
    82  }
    83  
    84  func (x *seqIterator) At() int { return x.c }
    85  
    86  func (x *seqIterator) Close() error { return nil }
    87  func (x *seqIterator) Err() error   { return nil }
    88  
    89  // Benchmark_AsyncBatchIterator-10    	   91417	     13353 ns/op	   17017 B/op	      10 allocs/op
    90  func Benchmark_AsyncBatchIterator(b *testing.B) {
    91  	b.ReportAllocs()
    92  	var n int
    93  	for i := 0; i < b.N; i++ {
    94  		x := NewAsyncBatchIterator[int, int](
    95  			newSeqIterator(1<<20),
    96  			1<<10,
    97  			func(i int) int { return i },
    98  			func([]int) {},
    99  		)
   100  		for x.Next() {
   101  			n += x.At()
   102  		}
   103  	}
   104  }
   105  
   106  // Benchmark_BufferedIterator-10    	      12	  99730976 ns/op	   10047 B/op	       8 allocs/op
   107  func Benchmark_BufferedIterator(b *testing.B) {
   108  	b.ReportAllocs()
   109  	var n int
   110  	for i := 0; i < b.N; i++ {
   111  		x := NewBufferedIterator[int](newSeqIterator(1<<20), 1<<10)
   112  		for x.Next() {
   113  			n += x.At()
   114  		}
   115  	}
   116  }