query_api.go

package queryapi

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/regexp"
	"github.com/munnerz/goautoneg"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/util/annotations"
	"github.com/prometheus/prometheus/util/httputil"
	v1 "github.com/prometheus/prometheus/web/api/v1"
	"github.com/weaveworks/common/httpgrpc"

	"github.com/cortexproject/cortex/pkg/distributed_execution"
	"github.com/cortexproject/cortex/pkg/engine"
	"github.com/cortexproject/cortex/pkg/querier"
	"github.com/cortexproject/cortex/pkg/querier/stats"
	"github.com/cortexproject/cortex/pkg/util"
	"github.com/cortexproject/cortex/pkg/util/api"
	"github.com/cortexproject/cortex/pkg/util/requestmeta"
)

type QueryAPI struct {
	queryable             storage.SampleAndChunkQueryable
	queryEngine           engine.QueryEngine
	now                   func() time.Time
	statsRenderer         v1.StatsRenderer
	logger                log.Logger
	codecs                []v1.Codec
	CORSOrigin            *regexp.Regexp
	timeoutClassification stats.PhaseTrackerConfig
}

func NewQueryAPI(
	qe engine.QueryEngine,
	q storage.SampleAndChunkQueryable,
	statsRenderer v1.StatsRenderer,
	logger log.Logger,
	codecs []v1.Codec,
	CORSOrigin *regexp.Regexp,
	timeoutClassification stats.PhaseTrackerConfig,
) *QueryAPI {
	return &QueryAPI{
		queryEngine:           qe,
		queryable:             q,
		statsRenderer:         statsRenderer,
		logger:                logger,
		codecs:                codecs,
		CORSOrigin:            CORSOrigin,
		now:                   time.Now,
		timeoutClassification: timeoutClassification,
	}
}

func (q *QueryAPI) RangeQueryHandler(r *http.Request) (result apiFuncResult) {
	// TODO(Sungjin1212): Change to emit basic error (not gRPC)
	start, err := util.ParseTime(r.FormValue("start"))
	if err != nil {
		return invalidParamError(err, "start")
	}
	end, err := util.ParseTime(r.FormValue("end"))
	if err != nil {
		return invalidParamError(err, "end")
	}
	if end < start {
		return invalidParamError(ErrEndBeforeStart, "end")
	}

	step, err := util.ParseDurationMs(r.FormValue("step"))
	if err != nil {
		return invalidParamError(err, "step")
	}

	if step <= 0 {
		return invalidParamError(ErrNegativeStep, "step")
	}

	// For safety, limit the number of returned points per timeseries.
	// This is sufficient for 60s resolution for a week or 1h resolution for a year.
	if (end-start)/step > 11000 {
		return apiFuncResult{nil, &apiError{errorBadData, ErrStepTooSmall}, nil, nil}
	}

	ctx := r.Context()

	// Always record query start time for phase tracking, regardless of feature flag.
	queryStats := stats.FromContext(ctx)
	queryStats.SetQueryStart(time.Now())

	if to := r.FormValue("timeout"); to != "" {
		var cancel context.CancelFunc
		timeout, err := util.ParseDurationMs(to)
		if err != nil {
			return invalidParamError(err, "timeout")
		}

		ctx, cancel = context.WithTimeout(ctx, convertMsToDuration(timeout))
		defer cancel()
	}

	cfg := q.timeoutClassification
	ctx, cancel, earlyResult := applyTimeoutClassification(ctx, queryStats, cfg)
	if cancel != nil {
		defer cancel()
	}
	if earlyResult != nil {
		return *earlyResult
	}

	opts, err := extractQueryOpts(r)
	if err != nil {
		return apiFuncResult{nil, &apiError{errorBadData, err}, nil, nil}
	}

	ctx = engine.AddEngineTypeToContext(ctx, r)
	ctx = querier.AddBlockStoreTypeToContext(ctx, r.Header.Get(querier.BlockStoreTypeHeader))

	var qry promql.Query
	startTime := convertMsToTime(start)
	endTime := convertMsToTime(end)
	stepDuration := convertMsToDuration(step)

	byteLP := []byte(r.PostFormValue("plan"))
	if len(byteLP) != 0 {
		logicalPlan, err := distributed_execution.Unmarshal(byteLP)
		if err != nil {
			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("invalid logical plan: %v", err)}, nil, nil}
		}
		qry, err = q.queryEngine.MakeRangeQueryFromPlan(ctx, q.queryable, opts, logicalPlan, startTime, endTime, stepDuration, r.FormValue("query"))
		if err != nil {
			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("failed to create range query from logical plan: %v", err)}, nil, nil}
		}
	} else { // if there is logical plan field is empty, fall back
		qry, err = q.queryEngine.NewRangeQuery(ctx, q.queryable, opts, r.FormValue("query"), startTime, endTime, stepDuration)
		if err != nil {
			return invalidParamError(httpgrpc.Errorf(http.StatusBadRequest, "%s", err.Error()), "query")
		}
	}

	// From now on, we must only return with a finalizer in the result (to
	// be called by the caller) or call qry.Close ourselves (which is
	// required in the case of a panic).
	defer func() {
		if result.finalizer == nil {
			qry.Close()
		}
	}()

	ctx = httputil.ContextFromRequest(ctx, r)

	res := qry.Exec(ctx)
	if res.Err != nil {
		// If the context was cancelled/timed out, apply timeout classification.
		if ctx.Err() != nil {
			if classified := q.classifyTimeout(ctx, queryStats, cfg, res.Warnings, qry.Close); classified != nil {
				return *classified
			}
		}

		return apiFuncResult{nil, returnAPIError(res.Err), res.Warnings, qry.Close}
	}

	warnings := res.Warnings
	qs := q.statsRenderer(ctx, qry.Stats(), r.FormValue("stats"))

	return apiFuncResult{&v1.QueryData{
		ResultType: res.Value.Type(),
		Result:     res.Value,
		Stats:      qs,
	}, nil, warnings, qry.Close}
}

func (q *QueryAPI) InstantQueryHandler(r *http.Request) (result apiFuncResult) {
	// TODO(Sungjin1212): Change to emit basic error (not gRPC)
	ts, err := util.ParseTimeParam(r, "time", q.now().Unix())
	if err != nil {
		return invalidParamError(err, "time")
	}

	ctx := r.Context()

	// Always record query start time for phase tracking, regardless of feature flag.
	queryStats := stats.FromContext(ctx)
	queryStats.SetQueryStart(time.Now())

	if to := r.FormValue("timeout"); to != "" {
		var cancel context.CancelFunc
		timeout, err := util.ParseDurationMs(to)
		if err != nil {
			return invalidParamError(err, "timeout")
		}

		ctx, cancel = context.WithDeadline(ctx, q.now().Add(convertMsToDuration(timeout)))
		defer cancel()
	}

	cfg := q.timeoutClassification
	ctx, cancel, earlyResult := applyTimeoutClassification(ctx, queryStats, cfg)
	if cancel != nil {
		defer cancel()
	}
	if earlyResult != nil {
		return *earlyResult
	}

	opts, err := extractQueryOpts(r)
	if err != nil {
		return apiFuncResult{nil, &apiError{errorBadData, err}, nil, nil}
	}

	ctx = engine.AddEngineTypeToContext(ctx, r)
	ctx = querier.AddBlockStoreTypeToContext(ctx, r.Header.Get(querier.BlockStoreTypeHeader))

	var qry promql.Query
	tsTime := convertMsToTime(ts)

	byteLP := []byte(r.PostFormValue("plan"))
	if len(byteLP) != 0 {
		logicalPlan, err := distributed_execution.Unmarshal(byteLP)
		if err != nil {
			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("invalid logical plan: %v", err)}, nil, nil}
		}
		qry, err = q.queryEngine.MakeInstantQueryFromPlan(ctx, q.queryable, opts, logicalPlan, tsTime, r.FormValue("query"))
		if err != nil {
			return apiFuncResult{nil, &apiError{errorInternal, fmt.Errorf("failed to create instant query from logical plan: %v", err)}, nil, nil}
		}
	} else { // if there is logical plan field is empty, fall back
		qry, err = q.queryEngine.NewInstantQuery(ctx, q.queryable, opts, r.FormValue("query"), tsTime)
		if err != nil {
			return invalidParamError(httpgrpc.Errorf(http.StatusBadRequest, "%s", err.Error()), "query")
		}
	}

	// From now on, we must only return with a finalizer in the result (to
	// be called by the caller) or call qry.Close ourselves (which is
	// required in the case of a panic).
	defer func() {
		if result.finalizer == nil {
			qry.Close()
		}
	}()

	ctx = httputil.ContextFromRequest(ctx, r)

	res := qry.Exec(ctx)
	if res.Err != nil {
		// If the context was cancelled/timed out, apply timeout classification.
		if ctx.Err() != nil {
			if classified := q.classifyTimeout(ctx, queryStats, cfg, res.Warnings, qry.Close); classified != nil {
				return *classified
			}
		}

		return apiFuncResult{nil, returnAPIError(res.Err), res.Warnings, qry.Close}
	}

	warnings := res.Warnings
	qs := q.statsRenderer(ctx, qry.Stats(), r.FormValue("stats"))

	return apiFuncResult{&v1.QueryData{
		ResultType: res.Value.Type(),
		Result:     res.Value,
		Stats:      qs,
	}, nil, warnings, qry.Close}
}

func (q *QueryAPI) Wrap(f apiFunc) http.HandlerFunc {
	hf := func(w http.ResponseWriter, r *http.Request) {
		httputil.SetCORS(w, q.CORSOrigin, r)

		result := f(r)
		if result.finalizer != nil {
			defer result.finalizer()
		}

		if result.err != nil {
			api.RespondFromGRPCError(q.logger, w, result.err.err)
			return
		}

		if result.data != nil {
			q.respond(w, r, result.data, result.warnings, r.FormValue("query"))
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}

	return CompressionHandler{
		Handler: http.HandlerFunc(hf),
	}.ServeHTTP
}

func (q *QueryAPI) respond(w http.ResponseWriter, req *http.Request, data any, warnings annotations.Annotations, query string) {
	warn, info := warnings.AsStrings(query, 10, 10)

	resp := &v1.Response{
		Status:   statusSuccess,
		Data:     data,
		Warnings: warn,
		Infos:    info,
	}

	codec, err := q.negotiateCodec(req, resp)
	if err != nil {
		api.RespondFromGRPCError(q.logger, w, httpgrpc.Errorf(http.StatusNotAcceptable, "%s", &apiError{errorNotAcceptable, err}))
		return
	}

	b, err := codec.Encode(resp)
	if err != nil {
		level.Error(q.logger).Log("error marshaling response", "url", req.URL, "err", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", codec.ContentType().String())
	w.Header().Set("X-Uncompressed-Length", strconv.Itoa(len(b)))
	w.WriteHeader(http.StatusOK)
	if n, err := w.Write(b); err != nil {
		level.Error(q.logger).Log("error writing response", "url", req.URL, "bytesWritten", n, "err", err)
	}
}

// applyTimeoutClassification creates a proactive context timeout that fires before
// the PromQL engine's own timeout, adjusted for queue wait time. Returns the
// (possibly wrapped) context, an optional cancel func, and an optional early-exit
// result when the entire timeout budget was already consumed in the queue.
func applyTimeoutClassification(ctx context.Context, queryStats *stats.QueryStats, cfg stats.PhaseTrackerConfig) (context.Context, context.CancelFunc, *apiFuncResult) {
	if !cfg.Enabled {
		return ctx, nil, nil
	}
	var queueWaitTime time.Duration
	queueJoin := queryStats.LoadQueueJoinTime()
	queueLeave := queryStats.LoadQueueLeaveTime()
	if !queueJoin.IsZero() && !queueLeave.IsZero() {
		queueWaitTime = queueLeave.Sub(queueJoin)
	}
	effectiveTimeout := cfg.TotalTimeout - queueWaitTime
	if effectiveTimeout <= 0 {
		return ctx, nil, &apiFuncResult{nil, &apiError{errorTimeout, httpgrpc.Errorf(http.StatusServiceUnavailable,
			"query timed out: query spent too long in scheduler queue")}, nil, nil}
	}
	ctx, cancel := context.WithTimeout(ctx, effectiveTimeout)
	return ctx, cancel, nil
}

// classifyTimeout inspects phase timings after a context cancellation/timeout
// and returns an apiFuncResult if the timeout should be converted to a 4XX user error.
// Returns nil if no conversion applies and the caller should use the default error path.
func (q *QueryAPI) classifyTimeout(ctx context.Context, queryStats *stats.QueryStats, cfg stats.PhaseTrackerConfig, warnings annotations.Annotations, closer func()) *apiFuncResult {
	if !stats.IsEnabled(ctx) {
		return nil
	}

	queryStats.SetQueryEnd(time.Now())

	decision := stats.DecideTimeoutResponse(queryStats, cfg)

	fetchTime := queryStats.LoadQueryStorageWallTime()
	queryEnd := queryStats.LoadQueryEnd()
	totalTime := queryEnd.Sub(queryStats.LoadQueryStart())
	evalTime := totalTime - fetchTime
	var queueWaitTime time.Duration
	queueJoin := queryStats.LoadQueueJoinTime()
	queueLeave := queryStats.LoadQueueLeaveTime()
	if !queueJoin.IsZero() && !queueLeave.IsZero() {
		queueWaitTime = queueLeave.Sub(queueJoin)
	}
	level.Warn(q.logger).Log(
		"msg", "query shard timed out with classification",
		"request_id", requestmeta.RequestIdFromContext(ctx),
		"query_start", queryStats.LoadQueryStart(),
		"query_end", queryEnd,
		"queue_wait_time", queueWaitTime,
		"query_storage_wall_time", fetchTime,
		"eval_time", evalTime,
		"total_time", totalTime,
		"wall_time", queryStats.LoadWallTime(),
		"response_series", queryStats.LoadResponseSeries(),
		"fetched_series_count", queryStats.LoadFetchedSeries(),
		"fetched_chunk_bytes", queryStats.LoadFetchedChunkBytes(),
		"fetched_data_bytes", queryStats.LoadFetchedDataBytes(),
		"fetched_samples_count", queryStats.LoadFetchedSamples(),
		"fetched_chunks_count", queryStats.LoadFetchedChunks(),
		"split_queries", queryStats.LoadSplitQueries(),
		"store_gateway_touched_postings_count", queryStats.LoadStoreGatewayTouchedPostings(),
		"store_gateway_touched_posting_bytes", queryStats.LoadStoreGatewayTouchedPostingBytes(),
		"scanned_samples", queryStats.LoadScannedSamples(),
		"peak_samples", queryStats.LoadPeakSamples(),
		"decision", decision,
		"conversion_enabled", cfg.Enabled,
	)

	if cfg.Enabled && decision == stats.UserError4XX {
		return &apiFuncResult{nil, &apiError{errorExec, httpgrpc.Errorf(http.StatusUnprocessableEntity,
			"query timed out: query spent too long in evaluation - consider simplifying your query")}, warnings, closer}
	}

	if cfg.Enabled {
		return &apiFuncResult{nil, &apiError{errorTimeout, httpgrpc.Errorf(http.StatusGatewayTimeout,
			"%s", ErrUpstreamRequestTimeout)}, warnings, closer}
	}

	return nil
}

func (q *QueryAPI) negotiateCodec(req *http.Request, resp *v1.Response) (v1.Codec, error) {
	for _, clause := range goautoneg.ParseAccept(req.Header.Get("Accept")) {
		for _, codec := range q.codecs {
			if codec.ContentType().Satisfies(clause) && codec.CanEncode(resp) {
				return codec, nil
			}
		}
	}

	defaultCodec := q.codecs[0]
	if !defaultCodec.CanEncode(resp) {
		return nil, fmt.Errorf("cannot encode response as %s", defaultCodec.ContentType())
	}

	return defaultCodec, nil
}