query tee proxy with support for comparison of responses (#2211)

cmd/querytee/main.go

@@ -0,0 +1,77 @@
+package main
+
+import (
+	"flag"
+	"os"
+
+	"github.com/go-kit/kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/weaveworks/common/logging"
+	"github.com/weaveworks/common/server"
+
+	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/tools/querytee"
+
+	"github.com/grafana/loki/pkg/loghttp"
+)
+
+type Config struct {
+	ServerMetricsPort int
+	LogLevel          logging.Level
+	ProxyConfig       querytee.ProxyConfig
+}
+
+func main() {
+	// Parse CLI flags.
+	cfg := Config{}
+	flag.IntVar(&cfg.ServerMetricsPort, "server.metrics-port", 9900, "The port where metrics are exposed.")
+	cfg.LogLevel.RegisterFlags(flag.CommandLine)
+	cfg.ProxyConfig.RegisterFlags(flag.CommandLine)
+	flag.Parse()
+
+	util.InitLogger(&server.Config{
+		LogLevel: cfg.LogLevel,
+	})
+
+	// Run the instrumentation server.
+	registry := prometheus.NewRegistry()
+	registry.MustRegister(prometheus.NewGoCollector())
+
+	i := querytee.NewInstrumentationServer(cfg.ServerMetricsPort, registry)
+	if err := i.Start(); err != nil {
+		level.Error(util.Logger).Log("msg", "Unable to start instrumentation server", "err", err.Error())
+		os.Exit(1)
+	}
+
+	// Run the proxy.
+	proxy, err := querytee.NewProxy(cfg.ProxyConfig, util.Logger, lokiReadRoutes(), registry)
+	if err != nil {
+		level.Error(util.Logger).Log("msg", "Unable to initialize the proxy", "err", err.Error())
+		os.Exit(1)
+	}
+
+	if err := proxy.Start(); err != nil {
+		level.Error(util.Logger).Log("msg", "Unable to start the proxy", "err", err.Error())
+		os.Exit(1)
+	}
+
+	proxy.Await()
+}
+
+func lokiReadRoutes() []querytee.Route {
+	samplesComparator := querytee.NewSamplesComparator()
+	samplesComparator.RegisterSamplesType(loghttp.ResultTypeStream, compareStreams)
+
+	return []querytee.Route{
+		{Path: "/loki/api/v1/query_range", RouteName: "api_v1_query_range", Methods: "GET", ResponseComparator: samplesComparator},
+		{Path: "/loki/api/v1/query", RouteName: "api_v1_query", Methods: "GET", ResponseComparator: samplesComparator},
+		{Path: "/loki/api/v1/label", RouteName: "api_v1_label", Methods: "GET", ResponseComparator: nil},
+		{Path: "/loki/api/v1/labels", RouteName: "api_v1_labels", Methods: "GET", ResponseComparator: nil},
+		{Path: "/loki/api/v1/label/{name}/values", RouteName: "api_v1_label_name_values", Methods: "GET", ResponseComparator: nil},
+		{Path: "/loki/api/v1/series", RouteName: "api_v1_series", Methods: "GET", ResponseComparator: nil},
+		{Path: "/api/prom/query", RouteName: "api_prom_query", Methods: "GET", ResponseComparator: samplesComparator},
+		{Path: "/api/prom/label", RouteName: "api_prom_label", Methods: "GET", ResponseComparator: nil},
+		{Path: "/api/prom/label/{name}/values", RouteName: "api_prom_label_name_values", Methods: "GET", ResponseComparator: nil},
+		{Path: "/api/prom/series", RouteName: "api_prom_series", Methods: "GET", ResponseComparator: nil},
+	}
+}

cmd/querytee/response_comparator.go

@@ -0,0 +1,69 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/go-kit/kit/log/level"
+
+	"github.com/grafana/loki/pkg/loghttp"
+)
+
+func compareStreams(expectedRaw, actualRaw json.RawMessage) error {
+	var expected, actual loghttp.Streams
+
+	err := json.Unmarshal(expectedRaw, &expected)
+	if err != nil {
+		return err
+	}
+	err = json.Unmarshal(actualRaw, &actual)
+	if err != nil {
+		return err
+	}
+
+	if len(expected) != len(actual) {
+		return fmt.Errorf("expected %d streams but got %d", len(expected), len(actual))
+	}
+
+	streamLabelsToIndexMap := make(map[string]int, len(expected))
+	for i, actualStream := range actual {
+		streamLabelsToIndexMap[actualStream.Labels.String()] = i
+	}
+
+	for _, expectedStream := range expected {
+		actualStreamIndex, ok := streamLabelsToIndexMap[expectedStream.Labels.String()]
+		if !ok {
+			return fmt.Errorf("expected stream %s missing from actual response", expectedStream.Labels)
+		}
+
+		actualStream := actual[actualStreamIndex]
+		expectedValuesLen := len(expectedStream.Entries)
+		actualValuesLen := len(actualStream.Entries)
+
+		if expectedValuesLen != actualValuesLen {
+			err := fmt.Errorf("expected %d values for stream %s but got %d", expectedValuesLen,
+				expectedStream.Labels, actualValuesLen)
+			if expectedValuesLen > 0 && actualValuesLen > 0 {
+				level.Error(util.Logger).Log("msg", err.Error(), "oldest-expected-ts", expectedStream.Entries[0].Timestamp.UnixNano(),
+					"newest-expected-ts", expectedStream.Entries[expectedValuesLen-1].Timestamp.UnixNano(),
+					"oldest-actual-ts", actualStream.Entries[0].Timestamp.UnixNano(), "newest-actual-ts", actualStream.Entries[actualValuesLen-1].Timestamp.UnixNano())
+			}
+			return err
+		}
+
+		for i, expectedSamplePair := range expectedStream.Entries {
+			actualSamplePair := actualStream.Entries[i]
+			if !expectedSamplePair.Timestamp.Equal(actualSamplePair.Timestamp) {
+				return fmt.Errorf("expected timestamp %v but got %v for stream %s", expectedSamplePair.Timestamp.UnixNano(),
+					actualSamplePair.Timestamp.UnixNano(), expectedStream.Labels)
+			}
+			if expectedSamplePair.Line != actualSamplePair.Line {
+				return fmt.Errorf("expected line %s for timestamp %v but got %s for stream %s", expectedSamplePair.Line,
+					expectedSamplePair.Timestamp.UnixNano(), actualSamplePair.Line, expectedStream.Labels)
+			}
+		}
+	}
+
+	return nil
+}

cmd/querytee/response_comparator_test.go

@@ -0,0 +1,102 @@
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCompareStreams(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		expected json.RawMessage
+		actual   json.RawMessage
+		err      error
+	}{
+		{
+			name:     "no streams",
+			expected: json.RawMessage(`[]`),
+			actual:   json.RawMessage(`[]`),
+		},
+		{
+			name: "no streams in actual response",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"]]}
+						]`),
+			actual: json.RawMessage(`[]`),
+			err:    errors.New("expected 1 streams but got 0"),
+		},
+		{
+			name: "extra stream in actual response",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"]]},
+							{"stream":{"foo1":"bar1"},"values":[["1","1"]]}
+						]`),
+			err: errors.New("expected 1 streams but got 2"),
+		},
+		{
+			name: "same number of streams but with different labels",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo1":"bar1"},"values":[["1","1"]]}
+						]`),
+			err: errors.New("expected stream {foo=\"bar\"} missing from actual response"),
+		},
+		{
+			name: "difference in number of samples",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","2"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"]]}
+						]`),
+			err: errors.New("expected 2 values for stream {foo=\"bar\"} but got 1"),
+		},
+		{
+			name: "difference in sample timestamp",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","2"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["3","2"]]}
+						]`),
+			err: errors.New("expected timestamp 2 but got 3 for stream {foo=\"bar\"}"),
+		},
+		{
+			name: "difference in sample value",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","2"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","3"]]}
+						]`),
+			err: errors.New("expected line 2 for timestamp 2 but got 3 for stream {foo=\"bar\"}"),
+		},
+		{
+			name: "correct samples",
+			expected: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","2"]]}
+						]`),
+			actual: json.RawMessage(`[
+							{"stream":{"foo":"bar"},"values":[["1","1"],["2","2"]]}
+						]`),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			err := compareStreams(tc.expected, tc.actual)
+			if tc.err == nil {
+				require.NoError(t, err)
+				return
+			}
+			require.Error(t, err)
+			require.Equal(t, tc.err.Error(), err.Error())
+		})
+	}
+}