Use Protobuf-based JSON output for sampling strategies

cmd/all-in-one/sampling_strategies_example.json

@@ -0,0 +1,43 @@
+{
+  "service_strategies": [
+    {
+      "service": "foo",
+      "type": "probabilistic",
+      "param": 0.8,
+      "operation_strategies": [
+        {
+          "operation": "op1",
+          "type": "probabilistic",
+          "param": 0.2
+        },
+        {
+          "operation": "op2",
+          "type": "probabilistic",
+          "param": 0.4
+        }
+      ]
+    },
+    {
+      "service": "bar",
+      "type": "ratelimiting",
+      "param": 5
+    }
+  ],
+  "default_strategy": {
+    "type": "probabilistic",
+    "param": 0.5,
+    "operation_strategies": [
+      {
+        "operation": "/health",
+        "type": "probabilistic",
+        "param": 0.0
+      },
+      {
+        "operation": "/metrics",
+        "type": "probabilistic",
+        "param": 0.0
+      }
+    ]
+  }
+}
+

model/converter/json/sampling.go

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package json
+
+import (
+	"strings"
+
+	"github.com/gogo/protobuf/jsonpb"
+
+	"github.com/jaegertracing/jaeger/proto-gen/api_v2"
+)
+
+// SamplingStrategyResponseToJSON defines the official way to generate
+// a JSON response from /sampling endpoints.
+func SamplingStrategyResponseToJSON(protoObj *api_v2.SamplingStrategyResponse) (string, error) {
+	// For backwards compatibility with Thrift-to-JSON encoding,
+	// we want the output to include "strategyType":"PROBABILISTIC" when appropriate.
+	// However, due to design oversight, the enum value for PROBABILISTIC is 0, so
+	// we need to set EmitDefaults=true. This in turns causes null fields to be emitted too,
+	// so we take care of them below.
+	jsonpbMarshaler := jsonpb.Marshaler{
+		EmitDefaults: true,
+	}
+
+	str, err := jsonpbMarshaler.MarshalToString(protoObj)
+	if err != nil {
+		return "", err
+	}
+
+	// Because we set EmitDefaults, jsonpb will also render null entries, so we remove them here.
+	str = strings.ReplaceAll(str, `"probabilisticSampling":null,`, "")
+	str = strings.ReplaceAll(str, `,"rateLimitingSampling":null`, "")
+	str = strings.ReplaceAll(str, `,"operationSampling":null`, "")
+
+	return str, nil
+}

model/converter/json/sampling_test.go

@@ -0,0 +1,92 @@
+// Copyright (c) 2023 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package json
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	thriftconv "github.com/jaegertracing/jaeger/model/converter/thrift/jaeger"
+	api_v1 "github.com/jaegertracing/jaeger/thrift-gen/sampling"
+)
+
+func TestSamplingStrategyResponseToJSON_Error(t *testing.T) {
+	_, err := SamplingStrategyResponseToJSON(nil)
+	assert.Error(t, err)
+}
+
+// TestSamplingStrategyResponseToJSON verifies that the function outputs
+// the same string as Thrift-based JSON marshaler.
+func TestSamplingStrategyResponseToJSON(t *testing.T) {
+	t.Run("probabilistic", func(t *testing.T) {
+		s := &api_v1.SamplingStrategyResponse{
+			StrategyType: api_v1.SamplingStrategyType_PROBABILISTIC,
+			ProbabilisticSampling: &api_v1.ProbabilisticSamplingStrategy{
+				SamplingRate: 0.42,
+			},
+		}
+		compareProtoAndThriftJSON(t, s)
+	})
+	t.Run("rateLimiting", func(t *testing.T) {
+		s := &api_v1.SamplingStrategyResponse{
+			StrategyType: api_v1.SamplingStrategyType_RATE_LIMITING,
+			RateLimitingSampling: &api_v1.RateLimitingSamplingStrategy{
+				MaxTracesPerSecond: 42,
+			},
+		}
+		compareProtoAndThriftJSON(t, s)
+	})
+	t.Run("operationSampling", func(t *testing.T) {
+		a := 11.2 // we need a pointer to value
+		s := &api_v1.SamplingStrategyResponse{
+			OperationSampling: &api_v1.PerOperationSamplingStrategies{
+				DefaultSamplingProbability:       0.42,
+				DefaultUpperBoundTracesPerSecond: &a,
+				DefaultLowerBoundTracesPerSecond: 2,
+				PerOperationStrategies: []*api_v1.OperationSamplingStrategy{
+					{
+						Operation: "foo",
+						ProbabilisticSampling: &api_v1.ProbabilisticSamplingStrategy{
+							SamplingRate: 0.42,
+						},
+					},
+					{
+						Operation: "bar",
+						ProbabilisticSampling: &api_v1.ProbabilisticSamplingStrategy{
+							SamplingRate: 0.42,
+						},
+					},
+				},
+			},
+		}
+		compareProtoAndThriftJSON(t, s)
+	})
+}
+
+func compareProtoAndThriftJSON(t *testing.T, thriftObj *api_v1.SamplingStrategyResponse) {
+	protoObj, err := thriftconv.ConvertSamplingResponseToDomain(thriftObj)
+	require.NoError(t, err)
+
+	s1, err := json.Marshal(thriftObj)
+	require.NoError(t, err)
+
+	s2, err := SamplingStrategyResponseToJSON(protoObj)
+	require.NoError(t, err)
+
+	assert.Equal(t, string(s1), s2)
+}

pkg/clientcfg/clientcfghttp/handler.go

@@ -25,8 +25,10 @@ import (
 	"github.com/gorilla/mux"
 
 	"github.com/jaegertracing/jaeger/cmd/agent/app/configmanager"
+	p2json "github.com/jaegertracing/jaeger/model/converter/json"
+	t2p "github.com/jaegertracing/jaeger/model/converter/thrift/jaeger"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
-	tSampling "github.com/jaegertracing/jaeger/thrift-gen/sampling"
+	api_v1 "github.com/jaegertracing/jaeger/thrift-gen/sampling"
 )
 
 const mimeTypeApplicationJSON = "application/json"
@@ -69,6 +71,9 @@ type HTTPHandler struct {
 		// Number of bad responses due to malformed thrift
 		BadThriftFailures metrics.Counter `metric:"http-server.errors" tags:"status=5xx,source=thrift"`
 
+		// Number of bad responses due to proto conversion
+		BadProtoFailures metrics.Counter `metric:"http-server.errors" tags:"status=5xx,source=proto"`
+
 		// Number of failed response writes from http server
 		WriteFailures metrics.Counter `metric:"http-server.errors" tags:"status=5xx,source=write"`
 	}
@@ -85,14 +90,19 @@ func NewHTTPHandler(params HTTPHandlerParams) *HTTPHandler {
 func (h *HTTPHandler) RegisterRoutes(router *mux.Router) {
 	prefix := h.params.BasePath
 	if h.params.LegacySamplingEndpoint {
-		router.HandleFunc(prefix+"/", func(w http.ResponseWriter, r *http.Request) {
-			h.serveSamplingHTTP(w, r, true /* thriftEnums092 */)
-		}).Methods(http.MethodGet)
+		router.HandleFunc(
+			prefix+"/",
+			func(w http.ResponseWriter, r *http.Request) {
+				h.serveSamplingHTTP(w, r, h.encodeThriftLegacy)
+			},
+		).Methods(http.MethodGet)
 	}
-
-	router.HandleFunc(prefix+"/sampling", func(w http.ResponseWriter, r *http.Request) {
-		h.serveSamplingHTTP(w, r, false /* thriftEnums092 */)
-	}).Methods(http.MethodGet)
+	router.HandleFunc(
+		prefix+"/sampling",
+		func(w http.ResponseWriter, r *http.Request) {
+			h.serveSamplingHTTP(w, r, h.encodeProto)
+		},
+	).Methods(http.MethodGet)
 
 	router.HandleFunc(prefix+"/baggageRestrictions", func(w http.ResponseWriter, r *http.Request) {
 		h.serveBaggageHTTP(w, r)
@@ -118,7 +128,11 @@ func (h *HTTPHandler) writeJSON(w http.ResponseWriter, json []byte) error {
 	return nil
 }
 
-func (h *HTTPHandler) serveSamplingHTTP(w http.ResponseWriter, r *http.Request, thriftEnums092 bool) {
+func (h *HTTPHandler) serveSamplingHTTP(
+	w http.ResponseWriter,
+	r *http.Request,
+	encoder func(strategy *api_v1.SamplingStrategyResponse) ([]byte, error),
+) {
 	service, err := h.serviceFromRequest(w, r)
 	if err != nil {
 		return
@@ -129,23 +143,40 @@ func (h *HTTPHandler) serveSamplingHTTP(w http.ResponseWriter, r *http.Request,
 		http.Error(w, fmt.Sprintf("collector error: %+v", err), http.StatusInternalServerError)
 		return
 	}
-	jsonBytes, err := json.Marshal(resp)
+	jsonBytes, err := encoder(resp)
 	if err != nil {
-		h.metrics.BadThriftFailures.Inc(1)
-		http.Error(w, "cannot marshall Thrift to JSON", http.StatusInternalServerError)
+		http.Error(w, "cannot marshall to JSON", http.StatusInternalServerError)
 		return
 	}
-	if thriftEnums092 {
-		jsonBytes = h.encodeThriftEnums092(jsonBytes)
-	}
 	if err = h.writeJSON(w, jsonBytes); err != nil {
 		return
 	}
-	if thriftEnums092 {
-		h.metrics.LegacySamplingRequestSuccess.Inc(1)
-	} else {
-		h.metrics.SamplingRequestSuccess.Inc(1)
+}
+
+func (h *HTTPHandler) encodeThriftLegacy(strategy *api_v1.SamplingStrategyResponse) ([]byte, error) {
+	jsonBytes, err := json.Marshal(strategy)
+	if err != nil {
+		h.metrics.BadThriftFailures.Inc(1)
+		return nil, err
+	}
+	jsonBytes = h.encodeThriftEnums092(jsonBytes)
+	h.metrics.LegacySamplingRequestSuccess.Inc(1)
+	return jsonBytes, nil
+}
+
+func (h *HTTPHandler) encodeProto(strategy *api_v1.SamplingStrategyResponse) ([]byte, error) {
+	pbs, err := t2p.ConvertSamplingResponseToDomain(strategy)
+	if err != nil {
+		h.metrics.BadProtoFailures.Inc(1)
+		return nil, fmt.Errorf("ConvertSamplingResponseToDomain failed: %w", err)
+	}
+	str, err := p2json.SamplingStrategyResponseToJSON(pbs)
+	if err != nil {
+		h.metrics.BadProtoFailures.Inc(1)
+		return nil, fmt.Errorf("SamplingStrategyResponseToJSON failed: %w", err)
 	}
+	h.metrics.SamplingRequestSuccess.Inc(1)
+	return []byte(str), nil
 }
 
 func (h *HTTPHandler) serveBaggageHTTP(w http.ResponseWriter, r *http.Request) {
@@ -167,9 +198,9 @@ func (h *HTTPHandler) serveBaggageHTTP(w http.ResponseWriter, r *http.Request) {
 	h.metrics.BaggageRequestSuccess.Inc(1)
 }
 
-var samplingStrategyTypes = []tSampling.SamplingStrategyType{
-	tSampling.SamplingStrategyType_PROBABILISTIC,
-	tSampling.SamplingStrategyType_RATE_LIMITING,
+var samplingStrategyTypes = []api_v1.SamplingStrategyType{
+	api_v1.SamplingStrategyType_PROBABILISTIC,
+	api_v1.SamplingStrategyType_RATE_LIMITING,
 }
 
 // Replace string enum values produced from Thrift 0.9.3 generated classes