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) {}