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

     1  package iter
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func Test_Tee_All(t *testing.T) {
    12  	type testCase struct {
    13  		name    string
    14  		items   int
    15  		bufSize int
    16  		iters   int
    17  	}
    18  	testCases := []testCase{
    19  		{
    20  			name:    "empty",
    21  			items:   0,
    22  			iters:   10,
    23  			bufSize: 512,
    24  		},
    25  		{
    26  			name:    "no iterators",
    27  			bufSize: 512,
    28  		},
    29  		{
    30  			name:    "single iterator",
    31  			items:   1000,
    32  			iters:   1,
    33  			bufSize: 512,
    34  		},
    35  		{
    36  			name:    "matches buffer size",
    37  			items:   512,
    38  			iters:   10,
    39  			bufSize: 512,
    40  		},
    41  		{
    42  			name:    "larger than buffer",
    43  			items:   1000,
    44  			iters:   10,
    45  			bufSize: 512,
    46  		},
    47  		{
    48  			name:    "less than buffer",
    49  			items:   7,
    50  			iters:   10,
    51  			bufSize: 512,
    52  		},
    53  	}
    54  	for _, tc := range testCases {
    55  		tc := tc
    56  		t.Run(tc.name, func(t *testing.T) {
    57  			var wg sync.WaitGroup
    58  			s := newSeqIterator(tc.items)
    59  			n := make([]int, tc.iters)
    60  			it := newTee[int](s, tc.iters, tc.bufSize)
    61  			for i, x := range it {
    62  				x := x
    63  				i := i
    64  				wg.Add(1)
    65  				go func() {
    66  					defer wg.Done()
    67  					for x.Next() {
    68  						n[i]++
    69  					}
    70  					assert.NoError(t, x.Close())
    71  					assert.NoError(t, x.Err())
    72  				}()
    73  			}
    74  			wg.Wait()
    75  			for i, v := range n {
    76  				assert.Equal(t, tc.items, v)
    77  				assert.False(t, it[i].Next())
    78  				assert.NoError(t, it[i].Close())
    79  				assert.NoError(t, it[i].Err())
    80  			}
    81  			assert.False(t, s.Next())
    82  			assert.NoError(t, s.Close())
    83  			assert.NoError(t, s.Err())
    84  		})
    85  	}
    86  }
    87  
    88  func Test_Tee_Lag(t *testing.T) {
    89  	const (
    90  		items   = 1000
    91  		iters   = 10
    92  		bufSize = 512
    93  	)
    94  	var wg sync.WaitGroup
    95  	s := newSeqIterator(items)
    96  	n := make([]int, iters)
    97  	it := newTee[int](s, iters, bufSize)
    98  	for i, x := range it {
    99  		x := x
   100  		i := i
   101  		wg.Add(1)
   102  		go func() {
   103  			defer wg.Done()
   104  			// Each iterator will consume i * 100 items.
   105  			for (n[i] < (i+1)*(items/iters)) && x.Next() {
   106  				n[i]++
   107  			}
   108  			assert.NoError(t, x.Close())
   109  			assert.NoError(t, x.Err())
   110  		}()
   111  	}
   112  	wg.Wait()
   113  	for i, v := range n {
   114  		assert.Equal(t, (i+1)*(items/iters), v)
   115  		assert.False(t, it[i].Next())
   116  		assert.NoError(t, it[i].Close())
   117  		assert.NoError(t, it[i].Err())
   118  	}
   119  	assert.False(t, s.Next())
   120  	assert.NoError(t, s.Close())
   121  	assert.NoError(t, s.Err())
   122  }
   123  
   124  func Test_Tee_BufferReuse(t *testing.T) {
   125  	const (
   126  		items   = 1 << 20
   127  		iters   = 2
   128  		bufSize = 512
   129  	)
   130  
   131  	var wg sync.WaitGroup
   132  	s := newSeqIterator(items)
   133  	n := make([]int, iters)
   134  	it := newTee[int](s, iters, bufSize)
   135  	for i, x := range it {
   136  		x := x
   137  		i := i
   138  		wg.Add(1)
   139  		go func() {
   140  			defer wg.Done()
   141  			var j int
   142  			for x.Next() {
   143  				n[i]++
   144  				j++
   145  				// Let others consume.
   146  				if j%4<<10 == 0 {
   147  					runtime.Gosched()
   148  				}
   149  			}
   150  			assert.NoError(t, x.Close())
   151  			assert.NoError(t, x.Err())
   152  		}()
   153  	}
   154  	wg.Wait()
   155  
   156  	for i, v := range n {
   157  		assert.Equal(t, items, v)
   158  		assert.False(t, it[i].Next())
   159  		assert.NoError(t, it[i].Close())
   160  		assert.NoError(t, it[i].Err())
   161  	}
   162  	assert.False(t, s.Next())
   163  	assert.NoError(t, s.Close())
   164  	assert.NoError(t, s.Err())
   165  
   166  	// Might be flaky.
   167  	// Typically, for the given test, the expected
   168  	// buffer capacity is within [10K:100K].
   169  	assert.Less(t, cap(it[0].(*tee[int]).s.v), 2*items)
   170  }