From 9a532300655e86aa4968aaae93385c939a0a740f Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 14 Sep 2023 14:54:12 -0700 Subject: [PATCH] Check context before notifying frontend and scheduler Signed-off-by: Justin Jung --- pkg/querier/worker/scheduler_processor.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pkg/querier/worker/scheduler_processor.go b/pkg/querier/worker/scheduler_processor.go index d4e549e1a9..7ee7419064 100644 --- a/pkg/querier/worker/scheduler_processor.go +++ b/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)