diff --git a/api.go b/api.go index 08df4889c..5f5b3807f 100644 --- a/api.go +++ b/api.go @@ -45,6 +45,7 @@ type API struct { holder *Holder cluster *cluster server *Server + tracker *queryTracker importWorkersWG sync.WaitGroup importWorkerPoolSize int @@ -95,6 +96,8 @@ func NewAPI(opts ...apiOption) (*API, error) { }() } + api.tracker = newQueryTracker() + return api, nil } @@ -130,6 +133,7 @@ func (api *API) validate(f apiMethod) error { func (api *API) Close() error { close(api.importWork) api.importWorkersWG.Wait() + api.tracker.Stop() return nil } @@ -146,6 +150,7 @@ func (api *API) Query(ctx context.Context, req *QueryRequest) (QueryResponse, er if err != nil { return QueryResponse{}, errors.Wrap(err, "parsing") } + defer api.tracker.Finish(api.tracker.Start(req.Query)) execOpts := &execOptions{ Remote: req.Remote, Profile: req.Profile, @@ -1701,6 +1706,13 @@ func (api *API) GetTransaction(ctx context.Context, id string, remote bool) (*Tr return t, err } +func (api *API) ActiveQueries(ctx context.Context) ([]ActiveQueryStatus, error) { + if err := api.validate(apiActiveQueries); err != nil { + return nil, errors.Wrap(err, "validating api method") + } + return api.tracker.ActiveQueries(), nil +} + type serverInfo struct { ShardWidth uint64 `json:"shardWidth"` Memory uint64 `json:"memory"` @@ -1752,6 +1764,7 @@ const ( apiFinishTransaction apiTransactions apiGetTransaction + apiActiveQueries ) var methodsCommon = map[apiMethod]struct{}{ @@ -1791,4 +1804,5 @@ var methodsNormal = map[apiMethod]struct{}{ apiFinishTransaction: {}, apiTransactions: {}, apiGetTransaction: {}, + apiActiveQueries: {}, } diff --git a/http/handler.go b/http/handler.go index 12773f3ee..f6e0ac001 100644 --- a/http/handler.go +++ b/http/handler.go @@ -23,6 +23,7 @@ import ( "fmt" "io" "math" + "mime" "net" "net/http" _ "net/http/pprof" // Imported for its side-effect of registering pprof endpoints with the server. @@ -358,6 +359,7 @@ func newRouter(handler *Handler) *mux.Router { router.HandleFunc("/transaction/{id}", handler.handlePostTransaction).Methods("POST").Name("PostTransaction") router.HandleFunc("/transaction/{id}/finish", handler.handlePostFinishTransaction).Methods("POST").Name("PostFinishTransaction") router.HandleFunc("/transactions", handler.handleGetTransactions).Methods("GET").Name("GetTransactions") + router.HandleFunc("/queries", handler.handleGetActiveQueries).Methods("GET").Name("GetActiveQueries") router.HandleFunc("/version", handler.handleGetVersion).Methods("GET").Name("GetVersion") // /internal endpoints are for internal use only; they may change at any time. @@ -475,9 +477,33 @@ func (h *Handler) handleHome(w http.ResponseWriter, _ *http.Request) { // headers are present, but none of them are "application/json" // (or any matching wildcard). Otherwise returns true. func validHeaderAcceptJSON(header http.Header) bool { + return validHeaderAcceptType(header, "application", "json") +} + +func validHeaderAcceptType(header http.Header, typ, subtyp string) bool { if v, found := header["Accept"]; found { for _, v := range v { - if v == "application/json" || v == "*/*" || v == "*/json" || v == "application/*" { + t, _, err := mime.ParseMediaType(v) + if err != nil { + switch err { + case mime.ErrInvalidMediaParameter: + // This is an optional feature, so we can keep going anyway. + default: + continue + } + } + spl := strings.SplitN(t, "/", 2) + if len(spl) < 2 { + continue + } + switch { + case spl[0] == typ && spl[1] == subtyp: + return true + case spl[0] == "*" && spl[1] == subtyp: + return true + case spl[0] == typ && spl[1] == "*": + return true + case spl[0] == "*" && spl[1] == "*": return true } } @@ -835,6 +861,52 @@ func (h *Handler) handlePostIndexAttrDiff(w http.ResponseWriter, r *http.Request } } +func (h *Handler) handleGetActiveQueries(w http.ResponseWriter, r *http.Request) { + var rtype string + switch { + case validHeaderAcceptType(r.Header, "text", "plain"): + rtype = "text/plain" + case validHeaderAcceptJSON(r.Header): + rtype = "application/json" + default: + http.Error(w, "no acceptable response type selected", http.StatusNotAcceptable) + return + } + queries, err := h.api.ActiveQueries(r.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", rtype) + switch rtype { + case "text/plain": + durations := make([]string, len(queries)) + for i, q := range queries { + durations[i] = q.Age.String() + } + var maxlen int + for _, l := range durations { + if len(l) > maxlen { + maxlen = len(l) + } + } + for i, q := range queries { + _, err := fmt.Fprintf(w, "%*s%q\n", -(maxlen + 2), durations[i], q.Query) + if err != nil { + h.logger.Printf("sending GetActiveQueries response: %s", err) + return + } + } + if _, err := w.Write([]byte{'\n'}); err != nil { + h.logger.Printf("sending GetActiveQueries response: %s", err) + } + case "application/json": + if err := json.NewEncoder(w).Encode(queries); err != nil { + h.logger.Printf("encoding GetActiveQueries response: %s", err) + } + } +} + type postIndexAttrDiffRequest struct { Blocks []pilosa.AttrBlock `json:"blocks"` } diff --git a/tracker.go b/tracker.go new file mode 100644 index 000000000..af5788137 --- /dev/null +++ b/tracker.go @@ -0,0 +1,125 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pilosa + +import ( + "sort" + "sync" + "time" +) + +type ActiveQueryStatus struct { + Query string `json:"query"` + Age time.Duration `json:"age"` +} + +type activeQuery struct { + query string + started time.Time +} + +type queryStatusUpdate struct { + q *activeQuery + end bool +} + +type queryTracker struct { + updates chan<- queryStatusUpdate + checks chan<- chan<- []*activeQuery + wg sync.WaitGroup + stop chan struct{} +} + +func newQueryTracker() *queryTracker { + done := make(chan struct{}) + updates := make(chan queryStatusUpdate, 128) + checks := make(chan chan<- []*activeQuery) + tracker := &queryTracker{ + updates: updates, + checks: checks, + stop: done, + } + tracker.wg.Add(1) + go func() { + defer tracker.wg.Done() + + activeQueries := make(map[*activeQuery]struct{}) + + for { + select { + case update := <-updates: + if update.end { + delete(activeQueries, update.q) + } else { + activeQueries[update.q] = struct{}{} + } + case check := <-checks: + out := make([]*activeQuery, len(activeQueries)) + i := 0 + for q := range activeQueries { + out[i] = q + i++ + } + check <- out + close(check) + case <-done: + return + } + } + }() + return tracker +} + +func (t *queryTracker) Start(query string) *activeQuery { + now := time.Now() + q := &activeQuery{query, now} + t.updates <- queryStatusUpdate{q, false} + return q +} + +func (t *queryTracker) Finish(q *activeQuery) { + t.updates <- queryStatusUpdate{q, true} +} + +func (t *queryTracker) ActiveQueries() []ActiveQueryStatus { + ch := make(chan []*activeQuery, 1) + t.checks <- ch + queries := <-ch + sort.Slice(queries, func(i, j int) bool { + switch { + case queries[i].started.Before(queries[j].started): + return true + case queries[i].started.After(queries[j].started): + return false + case queries[i].query < queries[j].query: + return true + case queries[i].query > queries[j].query: + return false + default: + return false + } + }) + now := time.Now() + out := make([]ActiveQueryStatus, len(queries)) + for i, v := range queries { + out[i] = ActiveQueryStatus{v.query, now.Sub(v.started)} + } + return out +} + +func (t *queryTracker) Stop() { + close(t.stop) + t.wg.Wait() +} diff --git a/tracker_test.go b/tracker_test.go new file mode 100644 index 000000000..80f6f84f2 --- /dev/null +++ b/tracker_test.go @@ -0,0 +1,42 @@ +// Copyright 2020 Pilosa Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pilosa + +import "testing" + +func TestQueryTracker(t *testing.T) { + tracker := newQueryTracker() + defer tracker.Stop() + + if queries := tracker.ActiveQueries(); len(queries) > 0 { + t.Fatalf("expected no active queries; found %v", queries) + } + + qs := tracker.Start("test query") + + var queries []ActiveQueryStatus + for len(queries) < 1 { + queries = tracker.ActiveQueries() + } + if len(queries) > 1 || queries[0].Query != "test query" { + t.Fatalf("unexpected queries: %v", queries) + } + + tracker.Finish(qs) + + for len(queries) > 0 { + queries = tracker.ActiveQueries() + } +}