Query top rows

cmd/statshouse-api/statshouse-api.go

@@ -406,6 +406,7 @@ func run(argv args, cfg *api.Config, vkuthPublicKeys map[string][]byte) error {
 	m.Path("/debug/pprof/trace").Methods("GET").HandlerFunc(api.HandleProfTrace)
 	m.Path("/debug/pprof/symbol").Methods("GET").HandlerFunc(api.HandleProfSymbol)
 	m.Path("/debug/500").Methods("GET").HandlerFunc(api.DumpInternalServerErrors)
+	m.Path("/debug/top/rows").Methods("GET").HandlerFunc(api.DumpQueryTopRows)
 	m.Router.PathPrefix("/").Methods("GET", "HEAD").HandlerFunc(f.HandleStatic)
 
 	h := http.Handler(m)

internal/api/handler.go

@@ -209,6 +209,27 @@ type (
 		errorsMu sync.RWMutex
 		errors   [32]error500
 		errorX   int
+
+		// TOP queries by memory usage
+		queryTop   []queryInfo
+		queryTopMu sync.Mutex
+	}
+
+	queryInfo struct {
+		queryArgs
+		queryStatistics
+	}
+
+	queryArgs struct {
+		promQL   string
+		timeFrom int64
+		timeTo   int64
+	}
+
+	queryStatistics struct {
+		rowCount int
+		colCount int
+		memUsage int
 	}
 
 	//easyjson:json
@@ -2458,6 +2479,11 @@ func (h *requestHandler) handleSeriesRequest(ctx context.Context, req seriesRequ
 	} else {
 		limit = req.numResults
 	}
+	h.queryArgs = queryArgs{
+		promQL:   req.promQL,
+		timeFrom: req.from.Unix(),
+		timeTo:   req.to.Unix(),
+	}
 	if opt.timeNow.IsZero() {
 		opt.timeNow = time.Now()
 	}
@@ -3384,3 +3410,23 @@ func (h *requestHandler) init(accessToken, version string) (err error) {
 	h.endpointStat.setAccessInfo(h.accessInfo)
 	return nil
 }
+
+func (h *requestHandler) reportQueryStatistics(stat queryStatistics) {
+	if stat.rowCount <= 0 {
+		return
+	}
+	h.queryTopMu.Lock()
+	defer h.queryTopMu.Unlock()
+	s := h.queryTop
+	s = append(s, queryInfo{
+		queryArgs:       h.queryArgs,
+		queryStatistics: stat,
+	})
+	for i := len(s); i > 1 && s[i-2].memUsage < s[i-1].memUsage; i-- {
+		s[i-2], s[i-1] = s[i-1], s[i-2]
+	}
+	if len(s) > 100 { // TOP 100
+		s = s[:100]
+	}
+	h.queryTop = s
+}

internal/api/http_router.go

@@ -33,6 +33,7 @@ type requestHandler struct {
 	requestVersion string
 	forceVersion   string
 	versionDice    func() string
+	queryArgs      queryArgs
 }
 
 type httpRequestHandler struct {
@@ -165,6 +166,23 @@ func DumpInternalServerErrors(r *httpRequestHandler) {
 	}
 }
 
+func DumpQueryTopRows(r *httpRequestHandler) {
+	w := r.Response()
+	if ok := r.accessInfo.insecureMode || r.accessInfo.bitAdmin; !ok {
+		w.WriteHeader(http.StatusForbidden)
+		return
+	}
+	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+	var s []queryInfo
+	r.queryTopMu.Lock()
+	s = make([]queryInfo, 0, len(r.queryTop))
+	s = append(s, r.queryTop...)
+	r.queryTopMu.Unlock()
+	for i := 0; i < len(s); i++ {
+		w.Write([]byte(fmt.Sprintf("%d:%dx%d:%s\n", s[i].memUsage, s[i].rowCount, s[i].colCount, s[i].promQL)))
+	}
+}
+
 func (r *httpRequestHandler) Response() http.ResponseWriter {
 	return &r.w
 }

internal/api/promql.go

@@ -643,6 +643,11 @@ func (h *requestHandler) QuerySeries(ctx context.Context, qry *promql.SeriesQuer
 		tagX = make(map[tsTags]int, len(tagX))
 	}
 	res.Meta.Total = len(res.Data)
+	h.reportQueryStatistics(queryStatistics{
+		rowCount: len(res.Data),
+		colCount: len(qry.Timescale.Time),
+		memUsage: 8 * len(res.Data) * len(qry.Timescale.Time),
+	})
 	succeeded = true // prevents deffered "cleanup"
 	return res, cleanup, nil
 }