Using jsoniter to encode promql responses

Signed-off-by: alanprot <[email protected]>

pkg/api/api.go

@@ -20,7 +20,6 @@
 package api
 
 import (
-	"encoding/json"
 	"fmt"
 	"net/http"
 	"os"
@@ -29,6 +28,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	jsoniter "github.com/json-iterator/go"
 	"github.com/klauspost/compress/gzhttp"
 	"github.com/opentracing/opentracing-go"
 	"github.com/prometheus/common/route"
@@ -66,6 +66,11 @@ var corsHeaders = map[string]string{
 	"Access-Control-Expose-Headers": "Date",
 }
 
+var (
+	// Let suse the same json codec used by upstream prometheus.
+	json = jsoniter.ConfigCompatibleWithStandardLibrary
+)
+
 // ThanosVersion contains build information about Thanos.
 type ThanosVersion struct {
 	Version   string `json:"version"`

pkg/api/api_test.go

@@ -17,7 +17,6 @@
 package api
 
 import (
-	"encoding/json"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -28,12 +27,48 @@ import (
 	"github.com/opentracing/opentracing-go"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/route"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
+	promApiV1 "github.com/prometheus/prometheus/web/api/v1"
 
 	"github.com/efficientgo/core/testutil"
 	extpromhttp "github.com/thanos-io/thanos/pkg/extprom/http"
 	"github.com/thanos-io/thanos/pkg/logging"
 )
 
+func TestMarshallMatrixNull(t *testing.T) {
+	var m promql.Matrix
+	result := response{
+		Status: StatusSuccess,
+		Data: promApiV1.QueryData{
+			ResultType: parser.ValueTypeMatrix,
+			Result:     m, // null
+		},
+	}
+
+	b1, err := json.Marshal(result)
+	if err != nil {
+		t.Fatalf("Error marshaling response body: %s", err)
+	}
+
+	exp := response{
+		Status: StatusSuccess,
+		Data: promApiV1.QueryData{
+			ResultType: parser.ValueTypeMatrix,
+			Result:     promql.Matrix{},
+		},
+	}
+
+	b2, err := json.Marshal(exp)
+	if err != nil {
+		t.Fatalf("Error marshaling response body: %s", err)
+	}
+
+	if !reflect.DeepEqual(b1, b2) {
+		t.Fatalf("Expected response \n%v\n but got \n%v\n", string(b1), string(b2))
+	}
+}
+
 func TestRespondSuccess(t *testing.T) {
 	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		Respond(w, "test", nil)

pkg/query/endpointset.go

@@ -396,6 +396,8 @@ func (e *EndpointSet) Update(ctx context.Context) {
 		wg.Add(1)
 		go func(spec *GRPCEndpointSpec) {
 			defer wg.Done()
+			ctx, cancel := context.WithTimeout(ctx, e.endpointInfoTimeout)
+			defer cancel()
 			newRef, err := e.newEndpointRef(spec)
 			if err != nil {
 				level.Warn(e.logger).Log("msg", "new endpoint creation failed", "err", err, "address", spec.Addr())