SearchTagValuesV2 use protobuf internally instead of json to reduce latency

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 * [FEATURE] TraceQL support for event scope and event:name intrinsic [#3708](https://github.com/grafana/tempo/pull/3708) (@stoewer)
 * [FEATURE] Flush blocks to storage from the metrics-generator [#3628](https://github.com/grafana/tempo/pull/3628) [#3691](https://github.com/grafana/tempo/pull/3691) (@mapno)
+* [ENHANCEMENT] Tag value lookup use protobuf internally for improved latency [#3731](https://github.com/grafana/tempo/pull/3731) (@mdisibio)
 * [ENHANCEMENT] Improve use of OTEL semantic conventions on the service graph [#3711](https://github.com/grafana/tempo/pull/3711) (@zalegrala)
 
 ## v2.5.0-rc.1

modules/frontend/combiner/common.go

@@ -70,8 +70,21 @@ func (c *genericCombiner[T]) AddResponse(r PipelineResponse) error {
 	}
 
 	partial := c.new() // instantiating directly requires additional type constraints. this seemed cleaner: https://stackoverflow.com/questions/69573113/how-can-i-instantiate-a-non-nil-pointer-of-type-argument-with-generic-go
-	if err := jsonpb.Unmarshal(res.Body, partial); err != nil {
-		return fmt.Errorf("error unmarshalling response body: %w", err)
+
+	switch res.Header.Get(api.HeaderContentType) {
+	case api.HeaderAcceptProtobuf:
+		b, err := io.ReadAll(res.Body)
+		if err != nil {
+			return fmt.Errorf("error reading response body")
+		}
+		if err := proto.Unmarshal(b, partial); err != nil {
+			return fmt.Errorf("error unmarshalling proto response body: %w", err)
+		}
+	default:
+		// Assume json
+		if err := jsonpb.Unmarshal(res.Body, partial); err != nil {
+			return fmt.Errorf("error unmarshalling response body: %w", err)
+		}
 	}
 
 	if err := c.combine(partial, c.current, r); err != nil {

modules/frontend/tag_sharder.go

@@ -299,6 +299,7 @@ func (s searchTagSharder) buildBackendRequests(ctx context.Context, tenantID str
 				errFn(err)
 				return
 			}
+			subR.Header.Set(api.HeaderAccept, api.HeaderAcceptProtobuf)
 			prepareRequestForQueriers(subR, tenantID, parent.URL.Path, subR.URL.Query())
 
 			key := cacheKey(keyPrefix, tenantID, hash, int64(searchReq.start()), int64(searchReq.end()), m, startPage, pages)
@@ -356,7 +357,7 @@ func (s searchTagSharder) buildIngesterRequest(ctx context.Context, tenantID str
 	if err != nil {
 		return nil, err
 	}
-
+	subR.Header.Set(api.HeaderAccept, api.HeaderAcceptProtobuf)
 	prepareRequestForQueriers(subR, tenantID, subR.URL.Path, subR.URL.Query())
 	return subR, nil
 }

modules/querier/http.go

@@ -312,44 +312,33 @@ func (q *Querier) SearchTagValuesV2Handler(w http.ResponseWriter, r *http.Reques
 	span, ctx := opentracing.StartSpanFromContext(ctx, "Querier.SearchTagValuesHandler")
 	defer span.Finish()
 
+	var resp *tempopb.SearchTagValuesV2Response
+	var err error
+
 	if !isSearchBlock {
-		req, err := api.ParseSearchTagValuesRequestV2(r)
+		var req *tempopb.SearchTagValuesRequest
+		req, err = api.ParseSearchTagValuesRequestV2(r)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-
-		resp, err := q.SearchTagValuesV2(ctx, req)
-		if err != nil {
-			handleError(w, err)
-			return
-		}
-
-		marshaller := &jsonpb.Marshaler{}
-		err = marshaller.Marshal(w, resp)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
+		resp, err = q.SearchTagValuesV2(ctx, req)
 	} else {
-		req, err := api.ParseSearchTagValuesBlockRequestV2(r)
+		var req *tempopb.SearchTagValuesBlockRequest
+		req, err = api.ParseSearchTagValuesBlockRequestV2(r)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-		resp, err := q.SearchTagValuesBlocksV2(ctx, req)
-		if err != nil {
-			handleError(w, err)
-			return
-		}
-		marshaller := &jsonpb.Marshaler{}
-		err = marshaller.Marshal(w, resp)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
+		resp, err = q.SearchTagValuesBlocksV2(ctx, req)
 	}
-	w.Header().Set(api.HeaderContentType, api.HeaderAcceptJSON)
+
+	if err != nil {
+		handleError(w, err)
+		return
+	}
+
+	writeFormattedContentForRequest(w, r, resp)
 }
 
 func (q *Querier) SpanMetricsSummaryHandler(w http.ResponseWriter, r *http.Request) {
@@ -482,3 +471,30 @@ func handleError(w http.ResponseWriter, err error) {
 
 	http.Error(w, err.Error(), http.StatusInternalServerError)
 }
+
+func writeFormattedContentForRequest(w http.ResponseWriter, req *http.Request, m proto.Message) {
+	switch req.Header.Get(api.HeaderAccept) {
+	case api.HeaderAcceptProtobuf:
+		b, err := proto.Marshal(m)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		_, err = w.Write(b)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set(api.HeaderContentType, api.HeaderAcceptProtobuf)
+	default:
+		marshaller := &jsonpb.Marshaler{}
+		err := marshaller.Marshal(w, m)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set(api.HeaderContentType, api.HeaderAcceptJSON)
+	}
+}