Report expensive queries in query frontend (#78)

Author: yuchen-db

cmd/thanos/query_frontend.go

@@ -147,6 +147,9 @@ func registerQueryFrontend(app *extkingpin.App) {
 	cmd.Flag("query-frontend.log-queries-longer-than", "Log queries that are slower than the specified duration. "+
 		"Set to 0 to disable. Set to < 0 to enable on all queries.").Default("0").DurationVar(&cfg.CortexHandlerConfig.LogQueriesLongerThan)
 
+	cmd.Flag("query-frontend.log-queries-more-expensive-than", "Log queries that fetch more data than the specified bytes. "+
+		"Set to 0 to disable. Otherwise must be a positive value. Default is 1GB").Default("1073741824").Uint64Var(&cfg.CortexHandlerConfig.LogQueriesMoreExpensiveThan)
+
 	cmd.Flag("query-frontend.log-failed-queries", "Log failed queries due to any reason").Default("true").BoolVar(&cfg.CortexHandlerConfig.LogFailedQueries)
 
 	cmd.Flag("failed-query-cache-capacity", "Capacity of cache for failed queries. 0 means this feature is disabled.").

internal/cortex/frontend/transport/handler.go

@@ -21,6 +21,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/thanos-io/thanos/internal/cortex/frontend/transport/utils"
+	"github.com/thanos-io/thanos/internal/cortex/querier/queryrange"
 	querier_stats "github.com/thanos-io/thanos/internal/cortex/querier/stats"
 	"github.com/thanos-io/thanos/internal/cortex/tenant"
 	"github.com/thanos-io/thanos/internal/cortex/util"
@@ -43,12 +44,13 @@ var (
 
 // HandlerConfig Config for a Handler.
 type HandlerConfig struct {
-	LogQueriesLongerThan     time.Duration `yaml:"log_queries_longer_than"`
-	MaxBodySize              int64         `yaml:"max_body_size"`
-	QueryStatsEnabled        bool          `yaml:"query_stats_enabled"`
-	LogFailedQueries         bool          `yaml:"log_failed_queries"`
-	FailedQueryCacheCapacity int           `yaml:"failed_query_cache_capacity"`
-	SlowQueryLogsUserHeader  string        `yaml:"slow_query_logs_user_header"`
+	LogQueriesLongerThan        time.Duration `yaml:"log_queries_longer_than"`
+	MaxBodySize                 int64         `yaml:"max_body_size"`
+	QueryStatsEnabled           bool          `yaml:"query_stats_enabled"`
+	LogFailedQueries            bool          `yaml:"log_failed_queries"`
+	FailedQueryCacheCapacity    int           `yaml:"failed_query_cache_capacity"`
+	SlowQueryLogsUserHeader     string        `yaml:"slow_query_logs_user_header"`
+	LogQueriesMoreExpensiveThan uint64        `yaml:"log_queries_more_expensive_than"`
 }
 
 // Handler accepts queries and forwards them to RoundTripper. It can log slow queries,
@@ -60,12 +62,13 @@ type Handler struct {
 	failedQueryCache *utils.FailedQueryCache
 
 	// Metrics.
-	querySeconds     *prometheus.CounterVec
-	querySeries      *prometheus.CounterVec
-	queryBytes       *prometheus.CounterVec
-	activeUsers      *util.ActiveUsersCleanupService
-	slowQueryCount   prometheus.Counter
-	failedQueryCount prometheus.Counter
+	querySeconds        *prometheus.CounterVec
+	querySeries         *prometheus.CounterVec
+	queryBytes          *prometheus.CounterVec
+	activeUsers         *util.ActiveUsersCleanupService
+	slowQueryCount      prometheus.Counter
+	failedQueryCount    prometheus.Counter
+	expensiveQueryCount prometheus.Counter
 }
 
 // NewHandler creates a new frontend handler.
@@ -119,6 +122,10 @@ func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logge
 		Name: "cortex_failed_query_total",
 		Help: "Total number of failed queries detected.",
 	})
+	h.expensiveQueryCount = promauto.With(reg).NewCounter(prometheus.CounterOpts{
+		Name: "cortex_expensive_query_total",
+		Help: "Total number of expensive queries detected.",
+	})
 	return h
 }
 
@@ -204,12 +211,17 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Check whether we should parse the query string.
 	shouldReportSlowQuery := f.cfg.LogQueriesLongerThan != 0 && queryResponseTime > f.cfg.LogQueriesLongerThan
-	if shouldReportSlowQuery || f.cfg.QueryStatsEnabled {
+	queryBytesFetched := queryrange.GetQueryBytesFetchedFromHeader(resp.Header)
+	shouldReportExpensiveQuery := f.cfg.LogQueriesMoreExpensiveThan != 0 && queryBytesFetched > f.cfg.LogQueriesMoreExpensiveThan
+	if shouldReportSlowQuery || shouldReportExpensiveQuery || f.cfg.QueryStatsEnabled {
 		queryString = f.parseRequestQueryString(r, buf)
 	}
 
 	if shouldReportSlowQuery {
-		f.reportSlowQuery(r, hs, queryString, queryResponseTime)
+		f.reportSlowQuery(r, hs, queryString, queryBytesFetched, queryResponseTime)
+	}
+	if shouldReportExpensiveQuery {
+		f.reportExpensiveQuery(r, queryString, queryBytesFetched, queryResponseTime)
 	}
 	if f.cfg.QueryStatsEnabled {
 		f.reportQueryStats(r, queryString, queryResponseTime, stats)
@@ -245,8 +257,37 @@ func (f *Handler) reportFailedQuery(r *http.Request, queryString url.Values, err
 	level.Error(util_log.WithContext(r.Context(), f.log)).Log(logMessage...)
 }
 
+func (f *Handler) reportExpensiveQuery(r *http.Request, queryString url.Values, queryBytesFetched uint64, queryResponseTime time.Duration) {
+	f.expensiveQueryCount.Inc()
+	// NOTE(GiedriusS): see https://github.com/grafana/grafana/pull/60301 for more info.
+	grafanaDashboardUID := "-"
+	if dashboardUID := r.Header.Get("X-Dashboard-Uid"); dashboardUID != "" {
+		grafanaDashboardUID = dashboardUID
+	}
+	grafanaPanelID := "-"
+	if panelID := r.Header.Get("X-Panel-Id"); panelID != "" {
+		grafanaPanelID = panelID
+	}
+	remoteUser, _, _ := r.BasicAuth()
+
+	logMessage := append([]interface{}{
+		"msg", "expensive query",
+		"method", r.Method,
+		"host", r.Host,
+		"path", r.URL.Path,
+		"remote_user", remoteUser,
+		"remote_addr", r.RemoteAddr,
+		"query_megabytes_fetched", queryBytesFetched / (1024 * 1024),
+		"grafana_dashboard_uid", grafanaDashboardUID,
+		"grafana_panel_id", grafanaPanelID,
+		"query_response_time", queryResponseTime.String(),
+	}, formatQueryString(queryString)...)
+
+	level.Error(util_log.WithContext(r.Context(), f.log)).Log(logMessage...)
+}
+
 // reportSlowQuery reports slow queries.
-func (f *Handler) reportSlowQuery(r *http.Request, responseHeaders http.Header, queryString url.Values, queryResponseTime time.Duration) {
+func (f *Handler) reportSlowQuery(r *http.Request, responseHeaders http.Header, queryString url.Values, queryBytesFetched uint64, queryResponseTime time.Duration) {
 	f.slowQueryCount.Inc()
 	// NOTE(GiedriusS): see https://github.com/grafana/grafana/pull/60301 for more info.
 	grafanaDashboardUID := "-"
@@ -278,6 +319,7 @@ func (f *Handler) reportSlowQuery(r *http.Request, responseHeaders http.Header,
 		"remote_user", remoteUser,
 		"remote_addr", r.RemoteAddr,
 		"time_taken", queryResponseTime.String(),
+		"query_megabytes_fetched", queryBytesFetched / (1024 * 1024),
 		"grafana_dashboard_uid", grafanaDashboardUID,
 		"grafana_panel_id", grafanaPanelID,
 		"trace_id", thanosTraceID,

internal/cortex/querier/queryrange/query_bytes_fetched.go

@@ -4,6 +4,7 @@
 package queryrange
 
 import (
+	"net/http"
 	"strconv"
 )
 
@@ -52,3 +53,21 @@ func QueryBytesFetchedHttpHeaderValue(response Response) []string {
 	}
 	return result
 }
+
+func getHeaderValue(hdr http.Header, key string) uint64 {
+	if val, ok := hdr[key]; ok {
+		if len(val) != 1 {
+			return 0
+		}
+		n, err := strconv.ParseUint(val[0], 10, 64)
+		if err != nil {
+			return 0
+		}
+		return n
+	}
+	return 0
+}
+
+func GetQueryBytesFetchedFromHeader(hdr http.Header) uint64 {
+	return getHeaderValue(hdr, QueryBytesFetchedHeaderName)
+}