github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/symdb/resolver_test.go (about)

     1  package symdb
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  	"sync/atomic"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/mock"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    14  )
    15  
    16  func Test_Resolver_Unreleased_Failed_Partition(t *testing.T) {
    17  	s := newBlockSuite(t, [][]string{{"testdata/profile.pb.gz"}})
    18  	defer s.teardown()
    19  	ctx, cancel := context.WithCancel(context.Background())
    20  	// Pass canceled context to make partition initialization to fail.
    21  	cancel()
    22  
    23  	r := NewResolver(ctx, s.reader)
    24  	r.AddSamples(0, s.indexed[0][0].Samples)
    25  	_, err := r.Tree()
    26  	require.ErrorIs(t, err, context.Canceled)
    27  	r.Release()
    28  
    29  	// This time we pass normal context.
    30  	r = NewResolver(context.Background(), s.reader)
    31  	r.AddSamples(0, s.indexed[0][0].Samples)
    32  	_, err = r.Tree()
    33  	require.NoError(t, err)
    34  	r.Release()
    35  }
    36  
    37  func Test_Resolver_Error_Propagation(t *testing.T) {
    38  	m := new(mockSymbolsReader)
    39  	m.On("Partition", mock.Anything, mock.Anything).Return(nil, io.EOF).Once()
    40  	r := NewResolver(context.Background(), m)
    41  	r.AddSamples(0, schemav1.Samples{})
    42  	_, err := r.Tree()
    43  	require.ErrorIs(t, err, io.EOF)
    44  	r.Release()
    45  }
    46  
    47  func Test_Resolver_Cancellation(t *testing.T) {
    48  	s := newBlockSuite(t, [][]string{{"testdata/profile.pb.gz"}})
    49  	defer s.teardown()
    50  	ctx, cancel := context.WithCancel(context.Background())
    51  	defer cancel()
    52  
    53  	const (
    54  		workers    = 10
    55  		iterations = 10
    56  		depth      = 5
    57  	)
    58  
    59  	var wg sync.WaitGroup
    60  	wg.Add(workers)
    61  
    62  	for i := 0; i < workers; i++ {
    63  		go func() {
    64  			defer wg.Done()
    65  			for j := 0; j < iterations; j++ {
    66  				for d := 0; d < depth; d++ {
    67  					func() {
    68  						r := NewResolver(contextCancelAfter(ctx, int64(d)), s.reader)
    69  						defer r.Release()
    70  						r.AddSamples(0, s.indexed[0][0].Samples)
    71  						_, _ = r.Tree()
    72  					}()
    73  				}
    74  			}
    75  		}()
    76  	}
    77  
    78  	wg.Wait()
    79  }
    80  
    81  type mockSymbolsReader struct{ mock.Mock }
    82  
    83  func (m *mockSymbolsReader) Partition(ctx context.Context, partition uint64) (PartitionReader, error) {
    84  	args := m.Called(ctx, partition)
    85  	r, _ := args.Get(0).(PartitionReader)
    86  	return r, args.Error(1)
    87  }
    88  
    89  type fakeContext struct {
    90  	context.Context
    91  	once sync.Once
    92  	ch   chan struct{}
    93  	c    atomic.Int64
    94  	n    int64
    95  }
    96  
    97  func contextCancelAfter(ctx context.Context, n int64) context.Context {
    98  	return &fakeContext{
    99  		ch:      make(chan struct{}),
   100  		Context: ctx,
   101  		n:       n,
   102  	}
   103  }
   104  
   105  func (f *fakeContext) Done() <-chan struct{} {
   106  	if f.c.Add(1) > f.n {
   107  		f.once.Do(func() {
   108  			close(f.ch)
   109  		})
   110  	}
   111  	return f.ch
   112  }
   113  
   114  func (f *fakeContext) Err() error {
   115  	if f.c.Load() > f.n {
   116  		return context.Canceled
   117  	}
   118  	return f.Context.Err()
   119  }