github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/worker/worker_test.go (about)

     1  package worker
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/grafana/dskit/services"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  	"google.golang.org/grpc"
    14  
    15  	"github.com/grafana/loki/pkg/util/test"
    16  )
    17  
    18  func TestResetConcurrency(t *testing.T) {
    19  	tests := []struct {
    20  		name                string
    21  		parallelism         int
    22  		maxConcurrent       int
    23  		numTargets          int
    24  		expectedConcurrency int
    25  	}{
    26  		{
    27  			name:                "Test create at least one processor per target",
    28  			parallelism:         0,
    29  			maxConcurrent:       0,
    30  			numTargets:          2,
    31  			expectedConcurrency: 2,
    32  		},
    33  		{
    34  			name:                "Test concurrency equal to parallelism * target when MatchMaxConcurrency is false",
    35  			parallelism:         4,
    36  			maxConcurrent:       0,
    37  			numTargets:          2,
    38  			expectedConcurrency: 8,
    39  		},
    40  		{
    41  			name:                "Test concurrency is correct when numTargets does not divide evenly into maxConcurrent",
    42  			parallelism:         1,
    43  			maxConcurrent:       7,
    44  			numTargets:          4,
    45  			expectedConcurrency: 7,
    46  		},
    47  		{
    48  			name:                "Test Total Parallelism dividing evenly",
    49  			parallelism:         1,
    50  			maxConcurrent:       6,
    51  			numTargets:          2,
    52  			expectedConcurrency: 6,
    53  		},
    54  		{
    55  			name:                "Test Total Parallelism at least one worker per target",
    56  			parallelism:         1,
    57  			maxConcurrent:       3,
    58  			numTargets:          6,
    59  			expectedConcurrency: 6,
    60  		},
    61  	}
    62  
    63  	for _, tt := range tests {
    64  		t.Run(tt.name, func(t *testing.T) {
    65  			cfg := Config{
    66  				Parallelism:           tt.parallelism,
    67  				MatchMaxConcurrency:   tt.maxConcurrent > 0,
    68  				MaxConcurrentRequests: tt.maxConcurrent,
    69  			}
    70  
    71  			w, err := newQuerierWorkerWithProcessor(cfg, NewMetrics(cfg, nil), log.NewNopLogger(), &mockProcessor{}, "", nil, nil)
    72  			require.NoError(t, err)
    73  			require.NoError(t, services.StartAndAwaitRunning(context.Background(), w))
    74  
    75  			for i := 0; i < tt.numTargets; i++ {
    76  				// gRPC connections are virtual... they don't actually try to connect until they are needed.
    77  				// This allows us to use dummy ports, and not get any errors.
    78  				w.AddressAdded(fmt.Sprintf("127.0.0.1:%d", i))
    79  			}
    80  
    81  			test.Poll(t, 250*time.Millisecond, tt.expectedConcurrency, func() interface{} {
    82  				return getConcurrentProcessors(w)
    83  			})
    84  
    85  			require.NoError(t, services.StopAndAwaitTerminated(context.Background(), w))
    86  			assert.Equal(t, 0, getConcurrentProcessors(w))
    87  		})
    88  	}
    89  }
    90  
    91  func getConcurrentProcessors(w *querierWorker) int {
    92  	result := 0
    93  	w.mu.Lock()
    94  	defer w.mu.Unlock()
    95  
    96  	for _, mgr := range w.managers {
    97  		result += int(mgr.currentProcessors.Load())
    98  	}
    99  
   100  	return result
   101  }
   102  
   103  type mockProcessor struct{}
   104  
   105  func (m mockProcessor) processQueriesOnSingleStream(ctx context.Context, _ *grpc.ClientConn, _ string) {
   106  	<-ctx.Done()
   107  }
   108  
   109  func (m mockProcessor) notifyShutdown(_ context.Context, _ *grpc.ClientConn, _ string) {}