Improve decoding of JSON responses.

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/Workiva/go-datastructures v1.0.53
 	github.com/aws/aws-lambda-go v1.17.0
 	github.com/bmatcuk/doublestar v1.2.2
+	github.com/buger/jsonparser v1.1.1
 	github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee
 	github.com/cespare/xxhash/v2 v2.1.1
 	github.com/containerd/fifo v0.0.0-20190226154929-a9fb20d87448 // indirect

go.sum

@@ -264,6 +264,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/bmizerany/pat v0.0.0-20170815010413-6226ea591a40/go.mod h1:8rLXio+WjiTceGBHIoTvn60HIbs7Hm7bcHjyrSqYB9c=
 github.com/boltdb/bolt v1.3.1/go.mod h1:clJnj/oiGkjum5o1McbSZDSLxVThjynRyGBgiAx27Ps=
+github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
+github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
 github.com/c-bata/go-prompt v0.2.2/go.mod h1:VzqtzE2ksDBcdln8G7mk2RX9QyGjH+OVqOCSiVIqS34=
 github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee h1:BnPxIde0gjtTnc9Er7cxvBk8DHLWhEux0SxayC8dP6I=
 github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee/go.mod h1:S/7n9copUssQ56c7aAgHqftWO4LTf4xY6CGWt8Bc+3M=

pkg/loghttp/entry.go

@@ -5,6 +5,7 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/buger/jsonparser"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/modern-go/reflect2"
 )
@@ -19,6 +20,36 @@ type Entry struct {
 	Line      string
 }
 
+func (e *Entry) UnmarshalJSON(data []byte) error {
+	var (
+		i          int
+		parseError error
+	)
+	_, err := jsonparser.ArrayEach(data, func(value []byte, _ jsonparser.ValueType, _ int, _ error) {
+		switch i {
+		case 0: // timestamp
+			ts, err := jsonparser.ParseInt(value)
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Timestamp = time.Unix(0, ts)
+		case 1: // value
+			v, err := jsonparser.ParseString(value)
+			if err != nil {
+				parseError = err
+				return
+			}
+			e.Line = v
+		}
+		i++
+	})
+	if parseError != nil {
+		return parseError
+	}
+	return err
+}
+
 type jsonExtension struct {
 	jsoniter.DummyExtension
 }

pkg/loghttp/labels.go

@@ -1,11 +1,12 @@
 package loghttp
 
 import (
-	"bytes"
 	"net/http"
 	"sort"
 	"strconv"
+	"strings"
 
+	"github.com/buger/jsonparser"
 	"github.com/gorilla/mux"
 
 	"github.com/grafana/loki/pkg/logproto"
@@ -20,14 +21,32 @@ type LabelResponse struct {
 // LabelSet is a key/value pair mapping of labels
 type LabelSet map[string]string
 
+func (l *LabelSet) UnmarshalJSON(data []byte) error {
+	if *l == nil {
+		*l = make(LabelSet)
+	}
+	return jsonparser.ObjectEach(data, func(key, val []byte, _ jsonparser.ValueType, _ int) error {
+		v, err := jsonparser.ParseString(val)
+		if err != nil {
+			return err
+		}
+		k, err := jsonparser.ParseString(key)
+		if err != nil {
+			return err
+		}
+		(*l)[k] = v
+		return nil
+	})
+}
+
 // Map coerces LabelSet into a map[string]string. This is useful for working with adapter types.
 func (l LabelSet) Map() map[string]string {
 	return l
 }
 
 // String implements the Stringer interface.  It returns a formatted/sorted set of label key/value pairs.
 func (l LabelSet) String() string {
-	var b bytes.Buffer
+	var b strings.Builder
 
 	keys := make([]string, 0, len(l))
 	for k := range l {

pkg/loghttp/query.go

@@ -7,6 +7,7 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/buger/jsonparser"
 	json "github.com/json-iterator/go"
 	"github.com/prometheus/common/model"
 
@@ -36,6 +37,22 @@ type QueryResponse struct {
 	Data   QueryResponseData `json:"data"`
 }
 
+func (q *QueryResponse) UnmarshalJSON(data []byte) error {
+	return jsonparser.ObjectEach(data, func(key, value []byte, dataType jsonparser.ValueType, offset int) error {
+		switch string(key) {
+		case "status":
+			q.Status = string(value)
+		case "data":
+			var responseData QueryResponseData
+			if err := responseData.UnmarshalJSON(value); err != nil {
+				return err
+			}
+			q.Data = responseData
+		}
+		return nil
+	})
+}
+
 // PushRequest models a log stream push
 type PushRequest struct {
 	Streams []*Stream `json:"streams"`
@@ -79,6 +96,22 @@ func (Matrix) Type() ResultType { return ResultTypeMatrix }
 // Streams is a slice of Stream
 type Streams []Stream
 
+func (ss *Streams) UnmarshalJSON(data []byte) error {
+	var parseError error
+	_, err := jsonparser.ArrayEach(data, func(value []byte, dataType jsonparser.ValueType, offset int, err error) {
+		var stream Stream
+		if err := stream.UnmarshalJSON(value); err != nil {
+			parseError = err
+			return
+		}
+		*ss = append(*ss, stream)
+	})
+	if parseError != nil {
+		return parseError
+	}
+	return err
+}
+
 func (s Streams) ToProto() []logproto.Stream {
 	if len(s) == 0 {
 		return nil
@@ -97,52 +130,90 @@ type Stream struct {
 	Entries []Entry  `json:"values"`
 }
 
-// UnmarshalJSON implements the json.Unmarshaler interface.
-func (q *QueryResponseData) UnmarshalJSON(data []byte) error {
-	unmarshal := struct {
-		Type       ResultType      `json:"resultType"`
-		Result     json.RawMessage `json:"result"`
-		Statistics stats.Result    `json:"stats"`
-	}{}
-
-	err := json.Unmarshal(data, &unmarshal)
-	if err != nil {
-		return err
+func (s *Stream) UnmarshalJSON(data []byte) error {
+	if s.Labels == nil {
+		s.Labels = LabelSet{}
 	}
-
-	var value ResultValue
-
-	// unmarshal results
-	switch unmarshal.Type {
-	case ResultTypeStream:
-		var s Streams
-		err = json.Unmarshal(unmarshal.Result, &s)
-		value = s
-	case ResultTypeMatrix:
-		var m Matrix
-		err = json.Unmarshal(unmarshal.Result, &m)
-		value = m
-	case ResultTypeVector:
-		var v Vector
-		err = json.Unmarshal(unmarshal.Result, &v)
-		value = v
-	case ResultTypeScalar:
-		var v Scalar
-		err = json.Unmarshal(unmarshal.Result, &v)
-		value = v
-	default:
-		return fmt.Errorf("unknown type: %s", unmarshal.Type)
+	if len(s.Entries) > 0 {
+		s.Entries = s.Entries[:0]
 	}
+	return jsonparser.ObjectEach(data, func(key, value []byte, ty jsonparser.ValueType, _ int) error {
+		switch string(key) {
+		case "stream":
+			if err := s.Labels.UnmarshalJSON(value); err != nil {
+				return err
+			}
+		case "values":
+			if ty == jsonparser.Null {
+				return nil
+			}
+			var parseError error
+			_, err := jsonparser.ArrayEach(value, func(value []byte, ty jsonparser.ValueType, _ int, _ error) {
+				if ty == jsonparser.Null {
+					return
+				}
+				var entry Entry
+				if err := entry.UnmarshalJSON(value); err != nil {
+					parseError = err
+					return
+				}
+				s.Entries = append(s.Entries, entry)
+			})
+			if parseError != nil {
+				return parseError
+			}
+			return err
+		}
+		return nil
+	})
+}
 
+// UnmarshalJSON implements the json.Unmarshaler interface.
+func (q *QueryResponseData) UnmarshalJSON(data []byte) error {
+	resultType, err := jsonparser.GetString(data, "resultType")
 	if err != nil {
 		return err
 	}
-
-	q.ResultType = unmarshal.Type
-	q.Result = value
-	q.Statistics = unmarshal.Statistics
-
-	return nil
+	q.ResultType = ResultType(resultType)
+
+	return jsonparser.ObjectEach(data, func(key, value []byte, dataType jsonparser.ValueType, _ int) error {
+		switch string(key) {
+		case "result":
+			switch q.ResultType {
+			case ResultTypeStream:
+				ss := Streams{}
+				if err := ss.UnmarshalJSON(value); err != nil {
+					return err
+				}
+				q.Result = ss
+			case ResultTypeMatrix:
+				var m Matrix
+				if err = json.Unmarshal(value, &m); err != nil {
+					return err
+				}
+				q.Result = m
+			case ResultTypeVector:
+				var v Vector
+				if err = json.Unmarshal(value, &v); err != nil {
+					return err
+				}
+				q.Result = v
+			case ResultTypeScalar:
+				var v Scalar
+				if err = json.Unmarshal(value, &v); err != nil {
+					return err
+				}
+				q.Result = v
+			default:
+				return fmt.Errorf("unknown type: %s", q.ResultType)
+			}
+		case "stats":
+			if err := json.Unmarshal(value, &q.Statistics); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
 }
 
 // Scalar is a single timestamp/float with no labels

pkg/loghttp/query_test.go

@@ -7,9 +7,11 @@ import (
 	"testing"
 	"time"
 
+	jsoniter "github.com/json-iterator/go"
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/pkg/logproto"
+	"github.com/grafana/loki/pkg/logqlmodel/stats"
 )
 
 func TestParseRangeQuery(t *testing.T) {
@@ -188,3 +190,79 @@ func TestStreams_ToProto(t *testing.T) {
 		})
 	}
 }
+
+func Test_QueryResponseUnmarshal(t *testing.T) {
+	for _, tt := range []QueryResponse{
+		{
+			Status: "ok",
+			Data: QueryResponseData{
+				ResultType: "streams",
+				Result:     Streams{},
+				Statistics: stats.Result{},
+			},
+		},
+		{
+			Status: "ok",
+			Data: QueryResponseData{
+				ResultType: "streams",
+				Result: Streams{
+					Stream{
+						Labels: LabelSet{"foo": "bar"},
+						Entries: []Entry{
+							{Timestamp: time.Unix(0, 1), Line: "1"},
+							{Timestamp: time.Unix(0, 2), Line: "2"},
+						},
+					},
+				},
+				Statistics: stats.Result{
+					Summary: stats.Summary{
+						BytesProcessedPerSecond: 1238,
+					},
+				},
+			},
+		},
+		{
+			Status: "ok",
+			Data: QueryResponseData{
+				ResultType: "streams",
+				Result: Streams{
+					Stream{
+						Labels: LabelSet{"foo": "bar"},
+						Entries: []Entry{
+							{Timestamp: time.Unix(0, 1), Line: "log line 1"},
+							{Timestamp: time.Unix(0, 2), Line: "some log line 2"},
+						},
+					},
+					Stream{
+						Labels: LabelSet{"bar": "buzz", "level": "err", "foo": "bar"},
+						Entries: []Entry{
+							{Timestamp: time.Unix(0, 1), Line: "1"},
+							{Timestamp: time.Unix(0, 2), Line: "2"},
+							{Timestamp: time.Unix(0, 2), Line: "2"},
+							{Timestamp: time.Unix(0, 2), Line: "2"},
+							{Timestamp: time.Unix(0, 2), Line: "2"},
+						},
+					},
+				},
+				Statistics: stats.Result{
+					Summary: stats.Summary{
+						BytesProcessedPerSecond: 1238,
+					},
+					Store: stats.Store{
+						TotalChunksRef: 1234123123,
+					},
+				},
+			},
+		},
+	} {
+		tt := tt
+		t.Run("", func(t *testing.T) {
+			b, err := jsoniter.Marshal(tt)
+			require.Nil(t, err)
+			var actual QueryResponse
+			err = actual.UnmarshalJSON(b)
+			require.Nil(t, err)
+			require.Equal(t, tt, actual)
+		})
+	}
+}