github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/worker/worker_test.go (about)

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