diff --git a/api.go b/api.go index 875d350d6..62002472e 100644 --- a/api.go +++ b/api.go @@ -146,6 +146,7 @@ func (api *API) Query(ctx context.Context, req *QueryRequest) (QueryResponse, er } execOpts := &execOptions{ Remote: req.Remote, + Profile: req.Profile, ExcludeRowAttrs: req.ExcludeRowAttrs, // NOTE: Kept for Pilosa 1.x compat. ExcludeColumns: req.ExcludeColumns, // NOTE: Kept for Pilosa 1.x compat. ColumnAttrs: req.ColumnAttrs, // NOTE: Kept for Pilosa 1.x compat. diff --git a/executor.go b/executor.go index 50ed7f89e..1cb5dcf4e 100644 --- a/executor.go +++ b/executor.go @@ -150,7 +150,16 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar return resp, err } } - + if opt.Profile { + var prof tracing.ProfiledSpan + prof, ctx = tracing.StartProfiledSpanFromContext(ctx, "Execute") + defer prof.Finish() + var ok bool + resp.Profile, ok = prof.(*tracing.Profile) + if !ok { + return resp, fmt.Errorf("profiling execution failed: %T is not tracing.Profile", prof) + } + } results, err := e.execute(ctx, index, q, shards, opt) if err != nil { return resp, err @@ -2959,6 +2968,7 @@ type mapResponse struct { // execOptions represents an execution context for a single Execute() call. type execOptions struct { Remote bool + Profile bool ExcludeRowAttrs bool ExcludeColumns bool ColumnAttrs bool diff --git a/handler.go b/handler.go index f42cf1993..6d1a8fc13 100644 --- a/handler.go +++ b/handler.go @@ -16,6 +16,8 @@ package pilosa import ( "encoding/json" + + "github.com/pilosa/pilosa/v2/tracing" ) // QueryRequest represent a request to process a query. @@ -42,6 +44,9 @@ type QueryRequest struct { // If true, indicates that query is part of a larger distributed query. // If false, this request is on the originating node. Remote bool + + // Should we profile this query? + Profile bool } // QueryResponse represent a response from a processed query. @@ -55,6 +60,9 @@ type QueryResponse struct { // Error during parsing or execution. Err error + + // Profiling data, if any + Profile *tracing.Profile } // MarshalJSON marshals QueryResponse into a JSON-encoded byte slice @@ -68,9 +76,11 @@ func (resp *QueryResponse) MarshalJSON() ([]byte, error) { return json.Marshal(struct { Results []interface{} `json:"results"` ColumnAttrSets []*ColumnAttrSet `json:"columnAttrs,omitempty"` + Profile *tracing.Profile `json:"profile,omitempty"` }{ Results: resp.Results, ColumnAttrSets: resp.ColumnAttrSets, + Profile: resp.Profile, }) } diff --git a/http/handler.go b/http/handler.go index db5d2a326..7e508a43d 100644 --- a/http/handler.go +++ b/http/handler.go @@ -187,7 +187,7 @@ func (h *Handler) populateValidators() { h.validators["DeleteField"] = queryValidationSpecRequired() h.validators["PostImport"] = queryValidationSpecRequired().Optional("clear", "ignoreKeyCheck") h.validators["PostImportRoaring"] = queryValidationSpecRequired().Optional("remote", "clear") - h.validators["PostQuery"] = queryValidationSpecRequired().Optional("shards", "columnAttrs", "excludeRowAttrs", "excludeColumns") + h.validators["PostQuery"] = queryValidationSpecRequired().Optional("shards", "columnAttrs", "excludeRowAttrs", "excludeColumns", "profile") h.validators["GetInfo"] = queryValidationSpecRequired() h.validators["RecalculateCaches"] = queryValidationSpecRequired() h.validators["GetSchema"] = queryValidationSpecRequired() @@ -1021,9 +1021,20 @@ func (h *Handler) readURLQueryRequest(r *http.Request) (*pilosa.QueryRequest, er return nil, errors.New("invalid shard argument") } + // Optional profiling + profile := false + profileString := q.Get("profile") + if profileString != "" { + profile, err = strconv.ParseBool(q.Get("profile")) + if err != nil { + return nil, fmt.Errorf("invalid profile argument: '%s' (should be true/false)", profileString) + } + } + return &pilosa.QueryRequest{ Query: query, Shards: shards, + Profile: profile, ColumnAttrs: q.Get("columnAttrs") == "true", ExcludeRowAttrs: q.Get("excludeRowAttrs") == "true", ExcludeColumns: q.Get("excludeColumns") == "true", diff --git a/tracing/tracing.go b/tracing/tracing.go index 3625166b0..a7d58a251 100644 --- a/tracing/tracing.go +++ b/tracing/tracing.go @@ -17,15 +17,51 @@ package tracing import ( "context" "net/http" + "time" ) // GlobalTracer is a single, global instance of Tracer. var GlobalTracer Tracer = NopTracer() -// StartSpanFromContext returnus a new child span and context from a given +// StartSpanFromContext returns a new child span and context from a given // context using the global tracer. func StartSpanFromContext(ctx context.Context, operationName string) (Span, context.Context) { - return GlobalTracer.StartSpanFromContext(ctx, operationName) + return startMaybeProfiledSpanFromContext(ctx, operationName, false) +} + +// StartProfiledSpanFromContext returns a new child span and context from a given +// context using the global tracer. +func StartProfiledSpanFromContext(ctx context.Context, operationName string) (ProfiledSpan, context.Context) { + span, ctx := startMaybeProfiledSpanFromContext(ctx, operationName, true) + return span.(ProfiledSpan), ctx +} + +// startMaybeProfiledSpanFromContext figures out whether it needs to make a profiling span. +func startMaybeProfiledSpanFromContext(ctx context.Context, operationName string, startProfiling bool) (Span, context.Context) { + var parent ProfiledSpan + makeProfile := startProfiling + // Parent context may or may not be profiled. + // If it is, we need to make a sub-profile. If it isn't, we need to make a new profile if + // startProfiling is set. + span, ok := spanFromContext(ctx) + if ok { + if parent, ok = span.(ProfiledSpan); ok { + makeProfile = true + } + } + // if we aren't making a profile, this is easy: + if !makeProfile { + return GlobalTracer.StartSpanFromContext(ctx, operationName) + } + newProf := &Profile{Name: operationName, Begin: time.Now(), KV: make(map[string]interface{})} + if parent != nil { + parent.AddChild(newProf) + } + inner, ctx := GlobalTracer.StartSpanFromContext(ctx, operationName) + newProf.inner = inner + // insert ourselves in the context + ctx = context.WithValue(ctx, arbitraryContextKey, newProf) + return newProf, ctx } // Tracer implements a generic distributed tracing interface. @@ -49,6 +85,49 @@ type Span interface { LogKV(alternatingKeyValues ...interface{}) } +// ProfiledSpan represents a span which profiles itself and its children. +type ProfiledSpan interface { + Span + Dump() interface{} // suitable for marshaling + AddChild(ProfiledSpan) +} + +// Profile represents the profiling data for a span. It also handles +// the bookkeeping for the underlying span, but this is unexported so +// it doesn't get unmarshaled. +type Profile struct { + inner Span + Name string + Begin, End time.Time `json:"-"` + Duration time.Duration + Children []ProfiledSpan `json:",omitempty"` + KV map[string]interface{} `json:",omitempty"` +} + +func (p *Profile) Finish() { + p.inner.Finish() + p.End = time.Now() + p.Duration = p.End.Sub(p.Begin) +} + +func (p *Profile) LogKV(alternatingKeyValues ...interface{}) { + for i := 0; i < len(alternatingKeyValues)-1; i += 2 { + if s, ok := alternatingKeyValues[i].(string); ok { + p.KV[s] = alternatingKeyValues[i+1] + } + } + p.inner.LogKV(alternatingKeyValues...) +} + +// returns something that json could probably marshal. +func (p *Profile) Dump() interface{} { + return p +} + +func (p *Profile) AddChild(child ProfiledSpan) { + p.Children = append(p.Children, child) +} + // NopTracer returns a tracer that doesn't do anything. func NopTracer() Tracer { return &nopTracer{} @@ -70,3 +149,12 @@ type nopSpan struct{} func (s *nopSpan) Finish() {} func (s *nopSpan) LogKV(alternatingKeyValues ...interface{}) {} + +type arbitraryContextKeyType int + +var arbitraryContextKey arbitraryContextKeyType + +func spanFromContext(ctx context.Context) (Span, bool) { + span, ok := ctx.Value(arbitraryContextKey).(Span) + return span, ok +}