add trivial execution-time profiling

What if you could ?profile=true on a query and get some
numbers back? That'd be really cool.

We already have tracing/spans, but right now, those only generate
any data if you have something set up for them to trace to. Add a
fancy wrapper that lets us generate our own tracing data, and dump
it into the request response, if ?profile=true.

We track wall-clock execution time, plus possible arbitrary K/V
pairs. Memory stats are not included, because obtaining them is
surprisingly expensive.

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.

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

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,
 	})
 }
 

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",

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
+}