Check context before notifying frontend and scheduler

Signed-off-by: Justin Jung <[email protected]>

pkg/querier/worker/scheduler_processor.go

@@ -158,6 +158,10 @@ func (sp *schedulerProcessor) querierLoop(c schedulerpb.SchedulerForQuerier_Quer
 			logger := util_log.WithContext(ctx, sp.log)
 			sp.runRequest(ctx, logger, request.QueryID, request.FrontendAddress, request.StatsEnabled, request.HttpRequest)
 
+			if err = ctx.Err(); err != nil {
+				return
+			}
+
 			// Report back to scheduler that processing of the query has finished.
 			if err := c.Send(&schedulerpb.QuerierToScheduler{}); err != nil {
 				level.Error(logger).Log("msg", "error notifying scheduler about finished query", "err", err, "addr", address)
@@ -187,6 +191,10 @@ func (sp *schedulerProcessor) runRequest(ctx context.Context, logger log.Logger,
 		level.Info(logger).Log("msg", "finished request", "status_code", response.Code, "response_size", len(response.GetBody()))
 	}
 
+	if err = ctx.Err(); err != nil {
+		return
+	}
+
 	// Ensure responses that are too big are not retried.
 	if len(response.Body) >= sp.maxMessageSize {
 		level.Error(logger).Log("msg", "response larger than max message size", "size", len(response.Body), "maxMessageSize", sp.maxMessageSize)