Allow replacing the stats struct before rendering JSON

This allows other implementations to inject their own statistics that
they're gathering in data linked from the context.Context. For example,
Cortex can inject its stats.Stats value under the `cortex` key.

Signed-off-by: Andrew Bloomgarden <blmgrdn@amazon.com>

util/stats/query_stats.go

@@ -110,15 +110,25 @@ type querySamples struct {
 	TotalQueryableSamples        int        `json:"totalQueryableSamples"`
 }
 
-// QueryStats currently only holding query timings.
+// BuiltinStats holds the statistics that Prometheus's core gathers.
-type QueryStats struct {
+type BuiltinStats struct {
 	Timings queryTimings  `json:"timings,omitempty"`
 	Samples *querySamples `json:"samples,omitempty"`
 }
 
+// QueryStats holds BuiltinStats and any other stats the particular
+// implementation wants to collect.
+type QueryStats interface {
+	Builtin() BuiltinStats
+}
+
+func (s *BuiltinStats) Builtin() BuiltinStats {
+	return *s
+}
+
 // NewQueryStats makes a QueryStats struct with all QueryTimings found in the
 // given TimerGroup.
-func NewQueryStats(s *Statistics) *QueryStats {
+func NewQueryStats(s *Statistics) QueryStats {
 	var (
 		qt      queryTimings
 		samples *querySamples
@@ -150,7 +160,7 @@ func NewQueryStats(s *Statistics) *QueryStats {
 		samples.TotalQueryableSamplesPerStep = sp.totalSamplesPerStepPoints()
 	}
 
-	qs := QueryStats{Timings: qt, Samples: samples}
+	qs := BuiltinStats{Timings: qt, Samples: samples}
 	return &qs
 }
 

web/api/v1/api.go

@@ -104,6 +104,15 @@ type RulesRetriever interface {
 	AlertingRules() []*rules.AlertingRule
 }
 
+type StatsRenderer func(context.Context, *stats.Statistics, string) stats.QueryStats
+
+func defaultStatsRenderer(ctx context.Context, s *stats.Statistics, param string) stats.QueryStats {
+	if param != "" {
+		return stats.NewQueryStats(s)
+	}
+	return nil
+}
+
 // PrometheusVersion contains build information about Prometheus.
 type PrometheusVersion struct {
 	Version   string `json:"version"`
@@ -177,15 +186,16 @@ type API struct {
 	ready                 func(http.HandlerFunc) http.HandlerFunc
 	globalURLOptions      GlobalURLOptions
 
-	db          TSDBAdminStats
+	db            TSDBAdminStats
-	dbDir       string
+	dbDir         string
-	enableAdmin bool
+	enableAdmin   bool
-	logger      log.Logger
+	logger        log.Logger
-	CORSOrigin  *regexp.Regexp
+	CORSOrigin    *regexp.Regexp
-	buildInfo   *PrometheusVersion
+	buildInfo     *PrometheusVersion
-	runtimeInfo func() (RuntimeInfo, error)
+	runtimeInfo   func() (RuntimeInfo, error)
-	gatherer    prometheus.Gatherer
+	gatherer      prometheus.Gatherer
-	isAgent     bool
+	isAgent       bool
+	statsRenderer StatsRenderer
 
 	remoteWriteHandler http.Handler
 	remoteReadHandler  http.Handler
@@ -222,6 +232,7 @@ func NewAPI(
 	buildInfo *PrometheusVersion,
 	gatherer prometheus.Gatherer,
 	registerer prometheus.Registerer,
+	statsRenderer StatsRenderer,
 ) *API {
 	a := &API{
 		QueryEngine:       qe,
@@ -246,10 +257,15 @@ func NewAPI(
 		buildInfo:        buildInfo,
 		gatherer:         gatherer,
 		isAgent:          isAgent,
+		statsRenderer:    defaultStatsRenderer,
 
 		remoteReadHandler: remote.NewReadHandler(logger, registerer, q, configFunc, remoteReadSampleLimit, remoteReadConcurrencyLimit, remoteReadMaxBytesInFrame),
 	}
 
+	if statsRenderer != nil {
+		a.statsRenderer = statsRenderer
+	}
+
 	if ap != nil {
 		a.remoteWriteHandler = remote.NewWriteHandler(logger, ap)
 	}
@@ -344,9 +360,9 @@ func (api *API) Register(r *route.Router) {
 }
 
 type queryData struct {
-	ResultType parser.ValueType  `json:"resultType"`
+	ResultType parser.ValueType `json:"resultType"`
-	Result     parser.Value      `json:"result"`
+	Result     parser.Value     `json:"result"`
-	Stats      *stats.QueryStats `json:"stats,omitempty"`
+	Stats      stats.QueryStats `json:"stats,omitempty"`
 }
 
 func invalidParamError(err error, parameter string) apiFuncResult {
@@ -399,10 +415,11 @@ func (api *API) query(r *http.Request) (result apiFuncResult) {
 	}
 
 	// Optional stats field in response if parameter "stats" is not empty.
-	var qs *stats.QueryStats
+	sr := api.statsRenderer
-	if r.FormValue("stats") != "" {
+	if sr == nil {
-		qs = stats.NewQueryStats(qry.Stats())
+		sr = defaultStatsRenderer
 	}
+	qs := sr(ctx, qry.Stats(), r.FormValue("stats"))
 
 	return apiFuncResult{&queryData{
 		ResultType: res.Value.Type(),
@@ -480,10 +497,11 @@ func (api *API) queryRange(r *http.Request) (result apiFuncResult) {
 	}
 
 	// Optional stats field in response if parameter "stats" is not empty.
-	var qs *stats.QueryStats
+	sr := api.statsRenderer
-	if r.FormValue("stats") != "" {
+	if sr == nil {
-		qs = stats.NewQueryStats(qry.Stats())
+		sr = defaultStatsRenderer
 	}
+	qs := sr(ctx, qry.Stats(), r.FormValue("stats"))
 
 	return apiFuncResult{&queryData{
 		ResultType: res.Value.Type(),

web/api/v1/api_test.go

@@ -30,6 +30,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/prometheus/util/stats"
+
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
@@ -522,6 +524,14 @@ func TestLabelNames(t *testing.T) {
 	}
 }
 
+type testStats struct {
+	Custom string `json:"custom"`
+}
+
+func (testStats) Builtin() (_ stats.BuiltinStats) {
+	return
+}
+
 func TestStats(t *testing.T) {
 	suite, err := promql.NewTest(t, ``)
 	require.NoError(t, err)
@@ -535,7 +545,7 @@ func TestStats(t *testing.T) {
 			return time.Unix(123, 0)
 		},
 	}
-	request := func(method string, param string) (*http.Request, error) {
+	request := func(method, param string) (*http.Request, error) {
 		u, err := url.Parse("http://example.com")
 		require.NoError(t, err)
 		q := u.Query()
@@ -555,6 +565,7 @@ func TestStats(t *testing.T) {
 
 	for _, tc := range []struct {
 		name     string
+		renderer StatsRenderer
 		param    string
 		expected func(*testing.T, interface{})
 	}{
@@ -574,7 +585,7 @@ func TestStats(t *testing.T) {
 				require.IsType(t, i, &queryData{})
 				qd := i.(*queryData)
 				require.NotNil(t, qd.Stats)
-				qs := qd.Stats
+				qs := qd.Stats.Builtin()
 				require.NotNil(t, qs.Timings)
 				require.Greater(t, qs.Timings.EvalTotalTime, float64(0))
 				require.NotNil(t, qs.Samples)
@@ -589,7 +600,7 @@ func TestStats(t *testing.T) {
 				require.IsType(t, i, &queryData{})
 				qd := i.(*queryData)
 				require.NotNil(t, qd.Stats)
-				qs := qd.Stats
+				qs := qd.Stats.Builtin()
 				require.NotNil(t, qs.Timings)
 				require.Greater(t, qs.Timings.EvalTotalTime, float64(0))
 				require.NotNil(t, qs.Samples)
@@ -597,8 +608,30 @@ func TestStats(t *testing.T) {
 				require.NotNil(t, qs.Samples.TotalQueryableSamplesPerStep)
 			},
 		},
+		{
+			name: "custom handler with known value",
+			renderer: func(ctx context.Context, s *stats.Statistics, p string) stats.QueryStats {
+				if p == "known" {
+					return testStats{"Custom Value"}
+				}
+				return nil
+			},
+			param: "known",
+			expected: func(t *testing.T, i interface{}) {
+				require.IsType(t, i, &queryData{})
+				qd := i.(*queryData)
+				require.NotNil(t, qd.Stats)
+				j, err := json.Marshal(qd.Stats)
+				require.NoError(t, err)
+				require.JSONEq(t, string(j), `{"custom":"Custom Value"}`)
+			},
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
+			before := api.statsRenderer
+			defer func() { api.statsRenderer = before }()
+			api.statsRenderer = tc.renderer
+
 			for _, method := range []string{http.MethodGet, http.MethodPost} {
 				ctx := context.Background()
 				req, err := request(method, tc.param)

web/web.go

@@ -340,6 +340,7 @@ func New(logger log.Logger, o *Options) *Handler {
 		h.versionInfo,
 		o.Gatherer,
 		o.Registerer,
+		nil,
 	)
 
 	if o.RoutePrefix != "/" {