feat(query): add a feature flag for query tracing (#19437)

cmd/influxd/main.go

@@ -7,7 +7,6 @@ import (
 	"os"
 	"time"
 
-	"github.com/influxdata/flux"
 	"github.com/influxdata/influxdb/v2"
 	"github.com/influxdata/influxdb/v2/cmd/influxd/generate"
 	"github.com/influxdata/influxdb/v2/cmd/influxd/launcher"
@@ -43,11 +42,6 @@ func main() {
 		},
 	)
 
-	// TODO: this should be removed in the future: https://github.com/influxdata/influxdb/issues/16220
-	if os.Getenv("QUERY_TRACING") == "1" {
-		flux.EnableExperimentalTracing()
-	}
-
 	if err := rootCmd.Execute(); err != nil {
 		os.Exit(1)
 	}

flags.yml

@@ -83,6 +83,13 @@
   contact: Query Team
   lifetime: temporary
 
+- name: Query Tracing
+  description: Turn on query tracing for queries that are sampled
+  key: queryTracing
+  default: false
+  contact: Query Team
+  lifetime: permanent
+
 - name: Simple Task Options Extraction
   description: Simplified task options extraction to avoid undefined functions when saving tasks
   key: simpleTaskOptionsExtraction

go.mod

@@ -47,7 +47,7 @@ require (
 	github.com/hashicorp/vault/api v1.0.2
 	github.com/imdario/mergo v0.3.9 // indirect
 	github.com/influxdata/cron v0.0.0-20191203200038-ded12750aac6
-	github.com/influxdata/flux v0.82.2
+	github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38
 	github.com/influxdata/httprouter v1.3.1-0.20191122104820-ee83e2772f69
 	github.com/influxdata/influxql v0.0.0-20180925231337-1cbfca8e56b6
 	github.com/influxdata/pkg-config v0.2.3

go.sum

@@ -207,6 +207,8 @@ github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LB
 github.com/go-stack/stack v1.8.0 h1:5SgMzNM5HxrEjV0ww2lTmX6E2Izsfxas4+YHWRs3Lsk=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-test/deep v1.0.1/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
+github.com/gofrs/uuid v3.3.0+incompatible h1:8K4tyRfvU1CYPgJsveYFQMhpFd/wXNM7iK6rR7UHz84=
+github.com/gofrs/uuid v3.3.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.0 h1:xU6/SpYbvkNYiptHJYEDRseDLvYE7wSqhYYNy0QSUzI=
 github.com/gogo/protobuf v1.2.0/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
@@ -350,8 +352,8 @@ github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NH
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/influxdata/cron v0.0.0-20191203200038-ded12750aac6 h1:OtjKkeWDjUbyMi82C7XXy7Tvm2LXMwiBBXyFIGNPaGA=
 github.com/influxdata/cron v0.0.0-20191203200038-ded12750aac6/go.mod h1:XabtPPW2qsCg0tl+kjaPU+cFS+CjQXEXbT1VJvHT4og=
-github.com/influxdata/flux v0.82.2 h1:VtoF8pbyoS+3QLQQmihSmV0Ly6g/A73x+3VBUp9t15g=
-github.com/influxdata/flux v0.82.2/go.mod h1:sAAIEgQTlTpsXCUQ49ymoRsKqraPzIb7F3paT72/lE0=
+github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38 h1:fdMJEXTneIsmEJOl6SGC92KJA/52cKkrq1HST+m+KdA=
+github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38/go.mod h1:9nhYCE1gdxLODGCb4e5+Aj4+gS42sNwdn0mP++ZroAQ=
 github.com/influxdata/httprouter v1.3.1-0.20191122104820-ee83e2772f69 h1:WQsmW0fXO4ZE/lFGIE84G6rIV5SJN3P3sjIXAP1a8eU=
 github.com/influxdata/httprouter v1.3.1-0.20191122104820-ee83e2772f69/go.mod h1:pwymjR6SrP3gD3pRj9RJwdl1j5s3doEEV8gS4X9qSzA=
 github.com/influxdata/influxql v0.0.0-20180925231337-1cbfca8e56b6 h1:CFx+pP90q/qg3spoiZjf8donE4WpAdjeJfPOcoNqkWo=

query/control/controller.go

@@ -32,6 +32,7 @@ import (
 	"github.com/influxdata/flux/runtime"
 	"github.com/influxdata/influxdb/v2"
 	"github.com/influxdata/influxdb/v2/kit/errors"
+	"github.com/influxdata/influxdb/v2/kit/feature"
 	"github.com/influxdata/influxdb/v2/kit/prom"
 	"github.com/influxdata/influxdb/v2/kit/tracing"
 	influxlogger "github.com/influxdata/influxdb/v2/logger"
@@ -208,6 +209,10 @@ func (c *Controller) Query(ctx context.Context, req *query.Request) (flux.Query,
 	for _, dep := range c.dependencies {
 		ctx = dep.Inject(ctx)
 	}
+	// Add per-transformation spans if the feature flag is set.
+	if feature.QueryTracing().Enabled(ctx) {
+		ctx = flux.WithExperimentalTracingEnabled(ctx)
+	}
 	q, err := c.query(ctx, req.Compiler)
 	if err != nil {
 		return q, err

query/control/controller_test.go

@@ -20,10 +20,14 @@ import (
 	"github.com/influxdata/flux/plan"
 	"github.com/influxdata/flux/plan/plantest"
 	"github.com/influxdata/flux/stdlib/universe"
+	"github.com/influxdata/influxdb/v2/kit/feature"
+	pmock "github.com/influxdata/influxdb/v2/mock"
 	"github.com/influxdata/influxdb/v2/query"
 	_ "github.com/influxdata/influxdb/v2/query/builtin"
 	"github.com/influxdata/influxdb/v2/query/control"
 	"github.com/influxdata/influxdb/v2/query/stdlib/influxdata/influxdb"
+	"github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/mocktracer"
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	"go.uber.org/zap/zaptest"
@@ -1289,6 +1293,102 @@ func TestController_ReserveMemoryWithoutExceedingMax(t *testing.T) {
 	validateUnusedMemory(t, reg, config)
 }
 
+func TestController_QueryTracing(t *testing.T) {
+	// temporarily install a mock tracer to see which spans are created.
+	oldTracer := opentracing.GlobalTracer()
+	defer opentracing.SetGlobalTracer(oldTracer)
+	mockTracer := mocktracer.New()
+	opentracing.SetGlobalTracer(mockTracer)
+
+	const memoryBytesQuotaPerQuery = 64
+	config := config
+	config.MemoryBytesQuotaPerQuery = memoryBytesQuotaPerQuery
+	ctrl, err := control.New(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer shutdown(t, ctrl)
+
+	flagger := pmock.NewFlagger(map[feature.Flag]interface{}{
+		feature.QueryTracing(): true,
+	})
+	plainCtx := context.Background()
+	withFlagger, err := feature.Annotate(plainCtx, flagger)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tcs := []struct {
+		name          string
+		ctx           context.Context
+		doNotWantSpan string
+		wantSpan      string
+	}{
+		{
+			name:          "feature flag off",
+			ctx:           plainCtx,
+			doNotWantSpan: "*executetest.AllocatingFromProcedureSpec",
+		},
+		{
+			name:     "feature flag on",
+			ctx:      withFlagger,
+			wantSpan: "*executetest.AllocatingFromProcedureSpec",
+		},
+	}
+	for _, tc := range tcs {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			mockTracer.Reset()
+
+			compiler := &mock.Compiler{
+				CompileFn: func(ctx context.Context) (flux.Program, error) {
+					// Return a program that will allocate one more byte than is allowed.
+					pts := plantest.PlanSpec{
+						Nodes: []plan.Node{
+							plan.CreatePhysicalNode("allocating-from-test", &executetest.AllocatingFromProcedureSpec{
+								ByteCount: 16,
+							}),
+							plan.CreatePhysicalNode("yield", &universe.YieldProcedureSpec{Name: "_result"}),
+						},
+						Edges: [][2]int{
+							{0, 1},
+						},
+						Resources: flux.ResourceManagement{
+							ConcurrencyQuota: 1,
+						},
+					}
+
+					ps := plantest.CreatePlanSpec(&pts)
+					prog := &lang.Program{
+						Logger:   zaptest.NewLogger(t),
+						PlanSpec: ps,
+					}
+
+					return prog, nil
+				},
+			}
+
+			// Depending on how the feature flag is set in the context,
+			// we may or may not do query tracing here.
+			q, err := ctrl.Query(tc.ctx, makeRequest(compiler))
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+
+			consumeResults(t, q)
+			gotSpans := make(map[string]struct{})
+			for _, span := range mockTracer.FinishedSpans() {
+				gotSpans[span.OperationName] = struct{}{}
+			}
+			if _, found := gotSpans[tc.doNotWantSpan]; tc.doNotWantSpan != "" && found {
+				t.Fatalf("did not want to find span %q but it was there", tc.doNotWantSpan)
+			}
+			if _, found := gotSpans[tc.wantSpan]; tc.wantSpan != "" && !found {
+				t.Fatalf("wanted to find span %q but it was not there", tc.wantSpan)
+			}
+		})
+	}
+}
+
 func consumeResults(tb testing.TB, q flux.Query) {
 	tb.Helper()
 	for res := range q.Results() {

query/promql/internal/promqltests/go.mod

@@ -10,7 +10,7 @@ require (
 	github.com/go-kit/kit v0.10.0 // indirect
 	github.com/google/go-cmp v0.5.0
 	github.com/hashicorp/go-rootcerts v1.0.2 // indirect
-	github.com/influxdata/flux v0.82.2
+	github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38
 	github.com/influxdata/influxdb/v2 v2.0.0-00010101000000-000000000000
 	github.com/influxdata/influxql v1.0.1 // indirect
 	github.com/influxdata/promql/v2 v2.12.0

query/promql/internal/promqltests/go.sum

@@ -235,6 +235,8 @@ github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LB
 github.com/go-stack/stack v1.8.0 h1:5SgMzNM5HxrEjV0ww2lTmX6E2Izsfxas4+YHWRs3Lsk=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-test/deep v1.0.1/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
+github.com/gofrs/uuid v3.3.0+incompatible h1:8K4tyRfvU1CYPgJsveYFQMhpFd/wXNM7iK6rR7UHz84=
+github.com/gofrs/uuid v3.3.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/googleapis v1.1.0/go.mod h1:gf4bu3Q80BeJ6H1S1vYPm8/ELATdvryBaNFGgqEef3s=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.0/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
@@ -389,8 +391,8 @@ github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NH
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/influxdata/cron v0.0.0-20191203200038-ded12750aac6 h1:OtjKkeWDjUbyMi82C7XXy7Tvm2LXMwiBBXyFIGNPaGA=
 github.com/influxdata/cron v0.0.0-20191203200038-ded12750aac6/go.mod h1:XabtPPW2qsCg0tl+kjaPU+cFS+CjQXEXbT1VJvHT4og=
-github.com/influxdata/flux v0.82.2 h1:VtoF8pbyoS+3QLQQmihSmV0Ly6g/A73x+3VBUp9t15g=
-github.com/influxdata/flux v0.82.2/go.mod h1:sAAIEgQTlTpsXCUQ49ymoRsKqraPzIb7F3paT72/lE0=
+github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38 h1:fdMJEXTneIsmEJOl6SGC92KJA/52cKkrq1HST+m+KdA=
+github.com/influxdata/flux v0.82.3-0.20200826163009-c9ec07c5aa38/go.mod h1:9nhYCE1gdxLODGCb4e5+Aj4+gS42sNwdn0mP++ZroAQ=
 github.com/influxdata/httprouter v1.3.1-0.20191122104820-ee83e2772f69 h1:WQsmW0fXO4ZE/lFGIE84G6rIV5SJN3P3sjIXAP1a8eU=
 github.com/influxdata/httprouter v1.3.1-0.20191122104820-ee83e2772f69/go.mod h1:pwymjR6SrP3gD3pRj9RJwdl1j5s3doEEV8gS4X9qSzA=
 github.com/influxdata/influxdb1-client v0.0.0-20191209144304-8bf82d3c094d/go.mod h1:qj24IKcXYK6Iy9ceXlo3Tc+vtHo9lIhSX5JddghvEPo=

query/stdlib/influxdata/influxdb/source.go

@@ -47,7 +47,7 @@ func (s *Source) Run(ctx context.Context) {
 	labelValues := s.m.getLabelValues(ctx, s.orgID, s.op)
 	start := time.Now()
 	var err error
-	if flux.IsExperimentalTracingEnabled() {
+	if flux.IsExperimentalTracingEnabled(ctx) {
 		span, ctxWithSpan := tracing.StartSpanFromContextWithOperationName(ctx, "source-"+s.op)
 		err = s.runner.run(ctxWithSpan)
 		span.Finish()