From bf380799db2c6eb2e117dd0fdf7a07f970d59b62 Mon Sep 17 00:00:00 2001
From: alanprot <alanprot@gmail.com>
Date: Wed, 5 Jun 2024 12:05:38 -0700
Subject: [PATCH] Using jsoniter to encode promql responses

Signed-off-by: alanprot <alanprot@gmail.com>
---
 pkg/api/api.go      |  7 ++++++-
 pkg/api/api_test.go | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pkg/api/api.go b/pkg/api/api.go
index 6c2ef037f6c..4c80ff2b424 100644
--- a/pkg/api/api.go
+++ b/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"`
diff --git a/pkg/api/api_test.go b/pkg/api/api_test.go
index 88c9a307728..e3efd8faa5d 100644
--- a/pkg/api/api_test.go
+++ b/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 marshalling 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 marshalling 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)