Fix cancel issue between Query Frontend and Query Schdeduler

pkg/lokifrontend/frontend/v2/frontend.go

@@ -220,15 +220,20 @@ enqueueAgain:
 		return nil, httpgrpc.Errorf(http.StatusInternalServerError, "failed to enqueue request")
 	}
 
+	fmt.Println("waiting on either cancel or response", "cancelCh", cancelCh)
 	select {
 	case <-ctx.Done():
 		if cancelCh != nil {
-			select {
-			case cancelCh <- freq.queryID:
-				// cancellation sent.
-			default:
-				// failed to cancel, ignore.
-			}
+			// NOTE(kavi): I think we don't need buffer channel.
+			// Let it block until it's workers receives it, We don't want to exist RoundTripGRPC without cancelling the downstream request started by frontend workers.
+			f.schedulerWorkers.sendRequestCancel(freq.queryID, cancelCh)
+
+			// select {
+			// case cancelCh <- freq.queryID:
+			// cancellation sent.
+			// default:
+			// 	// failed to cancel, ignore.
+			// }
 		}
 		return nil, ctx.Err()
 

pkg/lokifrontend/frontend/v2/frontend_scheduler_worker.go

@@ -2,6 +2,7 @@ package v2
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"sync"
 	"time"
@@ -146,6 +147,20 @@ func (f *frontendSchedulerWorkers) getWorkersCount() int {
 	return len(f.workers)
 }
 
+// sendRequestCancel sends cancellation to the "already scheduled" frontendRequest.
+// It will make sure the frontend worker that is responsible for the `reqID`
+// receives the cancel signal.
+func (f *frontendSchedulerWorkers) sendRequestCancel(reqID uint64, cancelCh chan<- uint64) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// There can be the case, where Frontend is running without any Frontend workers,
+	// sending cancel shoudn't block in those cases.
+	if len(f.workers) > 0 {
+		cancelCh <- reqID
+	}
+}
+
 func (f *frontendSchedulerWorkers) connectToScheduler(ctx context.Context, address string) (*grpc.ClientConn, error) {
 	// Because we only use single long-running method, it doesn't make sense to inject user ID, send over tracing or add metrics.
 	opts, err := f.cfg.GRPCClientConfig.DialOption(nil, nil)
@@ -276,6 +291,7 @@ func (w *frontendSchedulerWorker) schedulerLoop(loop schedulerpb.SchedulerForFro
 			return nil
 
 		case req := <-w.requestCh:
+			fmt.Println("Got the request", req.queryID)
 			err := loop.Send(&schedulerpb.FrontendToScheduler{
 				Type:            schedulerpb.ENQUEUE,
 				QueryID:         req.queryID,
@@ -325,6 +341,7 @@ func (w *frontendSchedulerWorker) schedulerLoop(loop schedulerpb.SchedulerForFro
 			}
 
 		case reqID := <-w.cancelCh:
+			fmt.Println("Got the cancel as well", reqID)
 			err := loop.Send(&schedulerpb.FrontendToScheduler{
 				Type:    schedulerpb.CANCEL,
 				QueryID: reqID,

pkg/lokifrontend/frontend/v2/frontend_test.go

@@ -176,6 +176,45 @@ func TestFrontendCancellation(t *testing.T) {
 	})
 }
 
+// Bug: If FrontendWorkers are busy, cancellation passed by Query frontend may not reach
+// all the frontend workers. (assumming we run multiple frontend workers)
+func TestFrontendWorkerCancellation(t *testing.T) {
+	f, ms := setupFrontend(t, nil)
+
+	// fmt.Println("workers count", f.schedulerWorkers.getWorkersCount(), "max-concurrency per worker", f.cfg.WorkerConcurrency)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	// send multiple requests > maxconcurrency of scheduler. So that it keeps all the frontend worker busy in serving requests.
+	reqCount := testFrontendWorkerConcurrency + 5
+	var wg sync.WaitGroup
+	for i := 0; i < reqCount; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			resp, err := f.RoundTripGRPC(user.InjectOrgID(ctx, "test"), &httpgrpc.HTTPRequest{})
+			require.EqualError(t, err, context.DeadlineExceeded.Error())
+			require.Nil(t, resp)
+		}()
+	}
+
+	wg.Wait()
+
+	// We wait a bit to make sure scheduler receives the cancellation request.
+	// 2 * reqCount because for every request, should also be corresponding cancel request
+	test.Poll(t, 5*time.Second, 2*reqCount, func() interface{} {
+		ms.mu.Lock()
+		defer ms.mu.Unlock()
+
+		return len(ms.msgs)
+	})
+
+	ms.checkWithLock(func() {
+		require.Equal(t, 2*reqCount, len(ms.msgs))
+	})
+}
+
 func TestFrontendFailedCancellation(t *testing.T) {
 	f, ms := setupFrontend(t, nil)