Frontend Refactor

cmd/tempo-cli/cmd-query-search.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"errors"
-	"fmt"
 	"io"
 	"net/http"
 	"path"
@@ -113,7 +112,7 @@ func (cmd *querySearchCmd) searchHTTP(req *tempopb.SearchRequest) error {
 		return err
 	}
 
-	fmt.Println(httpReq)
+	// fmt.Println(httpReq)
 	httpResp, err := http.DefaultClient.Do(httpReq)
 	if err != nil {
 		return err
@@ -126,7 +125,7 @@ func (cmd *querySearchCmd) searchHTTP(req *tempopb.SearchRequest) error {
 	}
 
 	if httpResp.StatusCode != http.StatusOK {
-		return errors.New("failed to query: " + string(body))
+		return errors.New("failed to query. body: " + string(body) + " status: " + httpResp.Status)
 	}
 
 	resp := &tempopb.SearchResponse{}

integration/e2e/multi_tenant_test.go

@@ -167,20 +167,6 @@ func TestMultiTenantSearch(t *testing.T) {
 			require.NoError(t, err)
 			require.NotEmpty(t, tagsValuesV2Resp.TagValues)
 
-			// assert tenant federation metrics
-			if tc.tenantSize > 1 {
-				for _, ta := range tenants {
-					matcher, err := labels.NewMatcher(labels.MatchEqual, "tenant", ta)
-					require.NoError(t, err)
-					// we should have 8 requests for each tenant
-					err = tempo.WaitSumMetricsWithOptions(e2e.Equals(8),
-						[]string{"tempo_query_frontend_multitenant_success_total"},
-						e2e.WithLabelMatchers(matcher),
-					)
-					require.NoError(t, err)
-				}
-			}
-
 			// check metrics for all routes
 			routeTable := []struct {
 				route    string
@@ -215,7 +201,7 @@ func TestMultiTenantSearch(t *testing.T) {
 
 			time.Sleep(2 * time.Second) // ensure that blocklist poller has built the blocklist
 			now := time.Now()
-			util.SearchStreamAndAssertTrace(t, grpcCtx, grpcClient, info, now.Add(-5*time.Minute).Unix(), now.Add(5*time.Minute).Unix())
+			util.SearchStreamAndAssertTrace(t, grpcCtx, grpcClient, info, now.Add(-10*time.Minute).Unix(), now.Add(10*time.Minute).Unix())
 			assertRequestCountMetric(t, tempo, "/tempopb.StreamingQuerier/Search", 1)
 
 			// test unsupported endpoint

modules/frontend/combiner/common.go

@@ -1,74 +1,104 @@
 package combiner
 
 import (
-	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"strings"
 	"sync"
 
+	"github.com/gogo/protobuf/jsonpb"
+	"github.com/gogo/protobuf/proto"
+	"github.com/gogo/status"
 	"github.com/grafana/tempo/pkg/api"
-	"github.com/grafana/tempo/pkg/tempopb"
+	"google.golang.org/grpc/codes"
 )
 
 type TResponse interface {
-	*tempopb.SearchResponse | *tempopb.SearchTagsResponse | *tempopb.SearchTagsV2Response | *tempopb.SearchTagValuesResponse | *tempopb.SearchTagValuesV2Response
+	proto.Message
 }
 
-type genericCombiner[R TResponse] struct {
+type genericCombiner[T TResponse] struct {
 	mu sync.Mutex
 
-	final R
+	current T // todo: state mgmt is mixed between the combiner and the various implementations. put it in one spot.
 
-	combine func(body io.ReadCloser, final R) error
-	result  func(R) (string, error)
+	new      func() T
+	combine  func(partial T, final T) error
+	finalize func(T) (T, error)
+	diff     func(T) (T, error) // currently only implemented by the search combiner. required for streaming
+	quit     func(T) bool
 
-	code          int
-	statusMessage string
-	err           error
+	//
+	httpStatusCode int
+	httpRespBody   string
 }
 
-func (c *genericCombiner[R]) AddRequest(res *http.Response, _ string) error {
+// AddResponse is used to add a http response to the combiner.
+func (c *genericCombiner[T]) AddResponse(res *http.Response, _ string) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	if c.shouldQuit() {
+	if res == nil {
 		return nil
 	}
 
-	c.code = res.StatusCode
+	// todo: reevaluate this. should the caller owner the lifecycle of the http.response body?
+	defer func() { _ = res.Body.Close() }()
+
+	if c.shouldQuit() {
+		return nil
+	}
 
+	c.httpStatusCode = res.StatusCode
 	if res.StatusCode != http.StatusOK {
 		bytesMsg, err := io.ReadAll(res.Body)
 		if err != nil {
 			return fmt.Errorf("error reading response body: %w", err)
 		}
-		c.statusMessage = string(bytesMsg)
-		return errors.New(c.statusMessage)
+		c.httpRespBody = string(bytesMsg)
+		c.httpStatusCode = res.StatusCode
+		// don't return error. the error path is reserved for unexpected errors.
+		// http pipeline errors should be returned through the final response. (Complete/TypedComplete/TypedDiff)
+		return nil
 	}
 
-	defer func() { _ = res.Body.Close() }()
-	if err := c.combine(res.Body, c.final); err != nil {
-		c.statusMessage = internalErrorMsg
-		c.err = fmt.Errorf("error unmarshalling response body: %w", err)
-		return c.err
+	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)
+	}
+
+	if err := c.combine(partial, c.current); err != nil {
+		c.httpRespBody = internalErrorMsg
+		return fmt.Errorf("error combining in combiner: %w", err)
 	}
 
 	return nil
 }
 
-func (c *genericCombiner[R]) Complete() (*http.Response, error) {
+// HTTPFinal, GRPCComplete, and GRPCDiff are all responsible for returning something
+// usable in grpc streaming/http response.
+func (c *genericCombiner[T]) HTTPFinal() (*http.Response, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	bodyString, err := c.result(c.final)
+	httpErr, _ := c.erroredResponse()
+	if httpErr != nil {
+		return httpErr, nil
+	}
+
+	final, err := c.finalize(c.current)
 	if err != nil {
 		return nil, err
 	}
 
+	bodyString, err := new(jsonpb.Marshaler).MarshalToString(final)
+	if err != nil {
+		return nil, fmt.Errorf("error marshalling response body: %w", err)
+	}
+
 	return &http.Response{
-		StatusCode: c.code,
+		StatusCode: c.httpStatusCode,
 		Header: http.Header{
 			api.HeaderContentType: {api.HeaderAcceptJSON},
 		},
@@ -77,11 +107,68 @@ func (c *genericCombiner[R]) Complete() (*http.Response, error) {
 	}, nil
 }
 
+func (c *genericCombiner[T]) GRPCFinal() (T, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	var empty T
+	_, grpcErr := c.erroredResponse()
+	if grpcErr != nil {
+		return empty, grpcErr
+	}
+
+	final, err := c.finalize(c.current)
+	if err != nil {
+		return empty, err
+	}
+
+	return final, nil
+}
+
+func (c *genericCombiner[T]) GRPCDiff() (T, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	var empty T
+	_, grpcErr := c.erroredResponse()
+	if grpcErr != nil {
+		return empty, grpcErr
+	}
+
+	diff, err := c.diff(c.current)
+	if err != nil {
+		return empty, err
+	}
+
+	return diff, nil
+}
+
+func (c *genericCombiner[T]) erroredResponse() (*http.Response, error) {
+	if c.httpStatusCode == http.StatusOK {
+		return nil, nil
+	}
+
+	// build grpc error and http response
+	var grpcErr error
+	if c.httpStatusCode/100 == 5 {
+		grpcErr = status.Error(codes.Internal, c.httpRespBody)
+	} else {
+		grpcErr = status.Error(codes.InvalidArgument, c.httpRespBody)
+	}
+	httpResp := &http.Response{
+		StatusCode: c.httpStatusCode,
+		Status:     http.StatusText(c.httpStatusCode),
+		Body:       io.NopCloser(strings.NewReader(c.httpRespBody)),
+	}
+
+	return httpResp, grpcErr
+}
+
 func (c *genericCombiner[R]) StatusCode() int {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	return c.code
+	return c.httpStatusCode
 }
 
 func (c *genericCombiner[R]) ShouldQuit() bool {
@@ -92,14 +179,18 @@ func (c *genericCombiner[R]) ShouldQuit() bool {
 }
 
 func (c *genericCombiner[R]) shouldQuit() bool {
-	if c.err != nil {
+	if c.httpStatusCode/100 == 5 { // Bail on 5xx
+		return true
+	}
+
+	if c.httpStatusCode/100 == 4 { // Bail on 4xx
 		return true
 	}
 
-	if c.code/100 == 5 { // Bail on 5xx
+	if c.quit != nil && c.quit(c.current) {
 		return true
 	}
 
-	// 2xx and 404 are OK
+	// 2xx
 	return false
 }