tests(launcher): add a launcher test for read window aggregate push down (#18412)

This adds a launcher test for the read window aggregate push down to
verify that it is done when a query is sent with the appropriate
pattern, the output is correct, and that the metric is incremented that
signals the push down happened.

Author: jsternberg

cmd/influxd/launcher/launcher.go

@@ -337,6 +337,7 @@ type Launcher struct {
 	secretStore     string
 
 	featureFlags map[string]string
+	flagger      feature.Flagger
 
 	// Query options.
 	concurrencyQuota                int
@@ -849,7 +850,7 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 		Addr: m.httpBindAddress,
 	}
 
-	flagger := feature.DefaultFlagger()
+	m.flagger = feature.DefaultFlagger()
 	if len(m.featureFlags) > 0 {
 		f, err := overrideflagger.Make(m.featureFlags, feature.ByKey)
 		if err != nil {
@@ -858,15 +859,15 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 			return err
 		}
 		m.log.Info("Running with feature flag overrides", zap.Any("overrides", m.featureFlags))
-		flagger = f
+		m.flagger = f
 	}
 
 	var sessionSvc platform.SessionService
 	{
 		sessionSvc = session.NewService(session.NewStorage(inmem.NewSessionStore()), userSvc, userResourceSvc, authSvc, time.Duration(m.sessionLength)*time.Minute)
 		sessionSvc = session.NewSessionMetrics(m.reg, sessionSvc)
 		sessionSvc = session.NewSessionLogger(m.log.With(zap.String("service", "session")), sessionSvc)
-		sessionSvc = session.NewServiceController(flagger, m.kvService, sessionSvc)
+		sessionSvc = session.NewServiceController(m.flagger, m.kvService, sessionSvc)
 	}
 
 	var labelSvc platform.LabelService
@@ -877,7 +878,7 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 			return err
 		}
 		ls := label.NewService(labelsStore)
-		labelSvc = label.NewLabelController(flagger, m.kvService, ls)
+		labelSvc = label.NewLabelController(m.flagger, m.kvService, ls)
 	}
 
 	m.apibackend = &http.APIBackend{
@@ -925,7 +926,7 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 		OrgLookupService:                m.kvService,
 		WriteEventRecorder:              infprom.NewEventRecorder("write"),
 		QueryEventRecorder:              infprom.NewEventRecorder("query"),
-		Flagger:                         flagger,
+		Flagger:                         m.flagger,
 		FlagsHandler:                    feature.NewFlagsHandler(kithttp.ErrorHandler(0), feature.ByKey),
 	}
 
@@ -993,7 +994,7 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 		labelSvc = label.NewLabelMetrics(m.reg, labelSvc)
 		newHandler := label.NewHTTPLabelHandler(m.log, labelSvc)
 
-		labelsHTTPServer = kithttp.NewFeatureHandler(feature.NewLabelPackage(), flagger, oldHandler, newHandler, newHandler.Prefix())
+		labelsHTTPServer = kithttp.NewFeatureHandler(feature.NewLabelPackage(), m.flagger, oldHandler, newHandler, newHandler.Prefix())
 	}
 
 	// feature flagging for new authorization service
@@ -1017,7 +1018,7 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 		authService = authorization.NewAuthLogger(authLogger, authService)
 
 		newHandler := authorization.NewHTTPAuthHandler(m.log, authService, ts, lookupSvc)
-		authHTTPServer = kithttp.NewFeatureHandler(feature.NewAuthPackage(), flagger, oldHandler, newHandler, newHandler.Prefix())
+		authHTTPServer = kithttp.NewFeatureHandler(feature.NewAuthPackage(), m.flagger, oldHandler, newHandler, newHandler.Prefix())
 	}
 
 	var oldSessionHandler nethttp.Handler
@@ -1033,8 +1034,8 @@ func (m *Launcher) run(ctx context.Context) (err error) {
 			http.WithResourceHandler(onboardHTTPServer),
 			http.WithResourceHandler(authHTTPServer),
 			http.WithResourceHandler(labelsHTTPServer),
-			http.WithResourceHandler(kithttp.NewFeatureHandler(feature.SessionService(), flagger, oldSessionHandler, sessionHTTPServer.SignInResourceHandler(), sessionHTTPServer.SignInResourceHandler().Prefix())),
-			http.WithResourceHandler(kithttp.NewFeatureHandler(feature.SessionService(), flagger, oldSessionHandler, sessionHTTPServer.SignOutResourceHandler(), sessionHTTPServer.SignOutResourceHandler().Prefix())),
+			http.WithResourceHandler(kithttp.NewFeatureHandler(feature.SessionService(), m.flagger, oldSessionHandler, sessionHTTPServer.SignInResourceHandler(), sessionHTTPServer.SignInResourceHandler().Prefix())),
+			http.WithResourceHandler(kithttp.NewFeatureHandler(feature.SessionService(), m.flagger, oldSessionHandler, sessionHTTPServer.SignOutResourceHandler(), sessionHTTPServer.SignOutResourceHandler().Prefix())),
 			http.WithResourceHandler(userHTTPServer.MeResourceHandler()),
 			http.WithResourceHandler(userHTTPServer.UserResourceHandler()),
 		)

cmd/influxd/launcher/launcher_helpers.go

@@ -20,10 +20,13 @@ import (
 	"github.com/influxdata/influxdb/v2/bolt"
 	influxdbcontext "github.com/influxdata/influxdb/v2/context"
 	"github.com/influxdata/influxdb/v2/http"
+	"github.com/influxdata/influxdb/v2/kit/feature"
 	"github.com/influxdata/influxdb/v2/mock"
 	"github.com/influxdata/influxdb/v2/pkg/httpc"
 	"github.com/influxdata/influxdb/v2/pkger"
 	"github.com/influxdata/influxdb/v2/query"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/prometheus/common/expfmt"
 )
 
 // TestLauncher is a test wrapper for launcher.Launcher.
@@ -224,6 +227,7 @@ func (tl *TestLauncher) MustExecuteQuery(query string) *QueryResults {
 // Callers of ExecuteQuery must call Done on the returned QueryResults.
 func (tl *TestLauncher) ExecuteQuery(q string) (*QueryResults, error) {
 	ctx := influxdbcontext.SetAuthorizer(context.Background(), mock.NewMockAuthorizer(true, nil))
+	ctx, _ = feature.Annotate(ctx, tl.flagger)
 	fq, err := tl.QueryController().Query(ctx, &query.Request{
 		Authorization:  tl.Auth,
 		OrganizationID: tl.Auth.OrgID,
@@ -408,6 +412,25 @@ func (tl *TestLauncher) HTTPClient(tb testing.TB) *httpc.Client {
 	return tl.httpClient
 }
 
+func (tl *TestLauncher) Metrics(tb testing.TB) (metrics map[string]*dto.MetricFamily) {
+	req := tl.HTTPClient(tb).
+		Get("/metrics").
+		RespFn(func(resp *nethttp.Response) error {
+			if resp.StatusCode != nethttp.StatusOK {
+				return fmt.Errorf("unexpected status code: %d %s", resp.StatusCode, resp.Status)
+			}
+			defer func() { _ = resp.Body.Close() }()
+
+			var parser expfmt.TextParser
+			metrics, _ = parser.TextToMetricFamilies(resp.Body)
+			return nil
+		})
+	if err := req.Do(context.Background()); err != nil {
+		tb.Fatal(err)
+	}
+	return metrics
+}
+
 // QueryResult wraps a single flux.Result with some helper methods.
 type QueryResult struct {
 	t *testing.T

cmd/influxd/launcher/query_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"html/template"
 	"io"
+	"io/ioutil"
 	"math/rand"
 	nethttp "net/http"
 	"strings"
@@ -15,6 +16,7 @@ import (
 	"time"
 
 	"github.com/influxdata/flux"
+	"github.com/influxdata/flux/csv"
 	"github.com/influxdata/flux/execute"
 	"github.com/influxdata/flux/execute/executetest"
 	"github.com/influxdata/flux/lang"
@@ -745,3 +747,95 @@ from(bucket: "%s")
 		t.Fatal(err)
 	}
 }
+
+func TestLauncher_Query_PushDownWindowAggregate(t *testing.T) {
+	l := launcher.RunTestLauncherOrFail(t, ctx,
+		"--feature-flags", "pushDownWindowAggregateCount=true")
+	l.SetupOrFail(t)
+	defer l.ShutdownOrFail(t, ctx)
+
+	l.WritePointsOrFail(t, `
+m0,k=k0 f=0i 0
+m0,k=k0 f=1i 1000000000
+m0,k=k0 f=2i 2000000000
+m0,k=k0 f=3i 3000000000
+m0,k=k0 f=4i 4000000000
+m0,k=k0 f=5i 5000000000
+m0,k=k0 f=6i 6000000000
+m0,k=k0 f=5i 7000000000
+m0,k=k0 f=0i 8000000000
+m0,k=k0 f=6i 9000000000
+m0,k=k0 f=6i 10000000000
+m0,k=k0 f=7i 11000000000
+m0,k=k0 f=5i 12000000000
+m0,k=k0 f=8i 13000000000
+m0,k=k0 f=9i 14000000000
+m0,k=k0 f=5i 15000000000
+`)
+
+	getReadRequestCount := func() uint64 {
+		const metricName = "query_influxdb_source_read_request_duration_seconds"
+		mf := l.Metrics(t)[metricName]
+		if mf != nil {
+			for _, m := range mf.Metric {
+				for _, label := range m.Label {
+					if label.GetName() == "op" && label.GetValue() == "readWindowAggregate" {
+						return m.Histogram.GetSampleCount()
+					}
+				}
+			}
+		}
+		return 0
+	}
+
+	for _, tt := range []struct {
+		name string
+		q    string
+		res  string
+	}{
+		{
+			name: "count",
+			q: `
+from(bucket: v.bucket)
+	|> range(start: 1970-01-01T00:00:00Z, stop: 1970-01-01T00:00:15Z)
+	|> aggregateWindow(every: 5s, fn: count)
+	|> drop(columns: ["_start", "_stop"])
+`,
+			res: `
+#datatype,string,long,long,string,string,string,dateTime:RFC3339
+#group,false,false,false,true,true,true,false
+#default,_result,,,,,,
+,result,table,_value,_field,_measurement,k,_time
+,,0,5,f,m0,k0,1970-01-01T00:00:05Z
+,,0,5,f,m0,k0,1970-01-01T00:00:10Z
+,,0,5,f,m0,k0,1970-01-01T00:00:15Z
+`,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			wantCount := getReadRequestCount() + 1
+
+			prelude := fmt.Sprintf("v = {bucket: \"%s\", timeRangeStart: 1970-01-01T00:00:00Z, timeRangeStop: 1970-01-01T00:00:15Z}", l.Bucket.Name)
+			queryStr := prelude + "\n" + tt.q
+			res := l.MustExecuteQuery(queryStr)
+			defer res.Done()
+			got := flux.NewSliceResultIterator(res.Results)
+			defer got.Release()
+
+			dec := csv.NewMultiResultDecoder(csv.ResultDecoderConfig{})
+			want, err := dec.Decode(ioutil.NopCloser(strings.NewReader(tt.res)))
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer want.Release()
+
+			if err := executetest.EqualResultIterators(want, got); err != nil {
+				t.Fatal(err)
+			}
+
+			if want, got := wantCount, getReadRequestCount(); want != got {
+				t.Fatalf("unexpected sample count -want/+got:\n\t- %d\n\t+ %d", want, got)
+			}
+		})
+	}
+}