Remove source of flakiness by using sync.Cond

Signed-off-by: Filip Petkovski <[email protected]>

pkg/queryfrontend/roundtrip_test.go

@@ -562,10 +562,6 @@ func TestRoundTripQueryCacheWithShardingMiddleware(t *testing.T) {
 			}
 
 			testutil.Equals(t, tc.expected, *res)
-
-			//if *res > tc.expected {
-			//	t.Fatalf("Expected to get less than or exactly %d requests, got %d", tc.expected, *res)
-			//}
 		}) {
 			break
 		}
@@ -844,6 +840,8 @@ func promqlResultsWithFailures(numFailures int) (*int, http.Handler) {
 		},
 	}
 
+	cond := sync.NewCond(&sync.Mutex{})
+	cond.L.Lock()
 	return &count, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		lock.Lock()
 		defer lock.Unlock()
@@ -852,16 +850,21 @@ func promqlResultsWithFailures(numFailures int) (*int, http.Handler) {
 		if numFailures > 0 {
 			numFailures--
 
-			// Allow other requests to execute
-			lock.Unlock()
-			<-time.After(200 * time.Millisecond)
-			lock.Lock()
-
+			// Wait for a successful request.
+			// Release the lock to allow other requests to execute.
+			if numFailures == 0 {
+				lock.Unlock()
+				cond.Wait()
+				lock.Lock()
+			}
 			w.WriteHeader(500)
+		} else {
+			cond.Broadcast()
 		}
 		if err := json.NewEncoder(w).Encode(q); err != nil {
 			panic(err)
 		}
+
 		count++
 	})
 }