Add tracing spans for rules API

pkg/api/query/v1.go

@@ -740,6 +740,15 @@ func NewRulesHandler(client rules.UnaryClient, enablePartialResponse bool) func(
 	}
 
 	return func(r *http.Request) (interface{}, []error, *api.ApiError) {
+		span, ctx := tracing.StartSpan(r.Context(), "receive_http_request")
+		defer span.Finish()
+
+		var (
+			groups   *rulespb.RuleGroups
+			warnings storage.Warnings
+			err      error
+		)
+
 		typeParam := r.URL.Query().Get("type")
 		typ, ok := rulespb.RulesRequest_Type_value[strings.ToUpper(typeParam)]
 		if !ok {
@@ -754,7 +763,9 @@ func NewRulesHandler(client rules.UnaryClient, enablePartialResponse bool) func(
 			Type:                    rulespb.RulesRequest_Type(typ),
 			PartialResponseStrategy: ps,
 		}
-		groups, warnings, err := client.Rules(r.Context(), req)
+		tracing.DoInSpan(ctx, "retrieve_rules", func(ctx context.Context) {
+			groups, warnings, err = client.Rules(ctx, req)
+		})
 		if err != nil {
 			return nil, nil, &api.ApiError{Typ: api.ErrorInternal, Err: errors.Errorf("error retrieving rules: %v", err)}
 		}

pkg/rules/manager.go

@@ -28,6 +28,7 @@ import (
 	"github.com/thanos-io/thanos/pkg/rules/rulespb"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
+	"github.com/thanos-io/thanos/pkg/tracing"
 )
 
 const tmpRuleDir = ".tmp-rules"
@@ -390,7 +391,7 @@ func (m *Manager) Update(evalInterval time.Duration, files []string) error {
 }
 
 // Rules returns specified rules from manager. This is used by gRPC and locally for HTTP and UI purposes.
-func (m *Manager) Rules(r *rulespb.RulesRequest, s rulespb.Rules_RulesServer) error {
+func (m *Manager) Rules(r *rulespb.RulesRequest, s rulespb.Rules_RulesServer) (err error) {
 	groups := m.protoRuleGroups()
 
 	pgs := make([]*rulespb.RuleGroup, 0, len(groups))
@@ -419,7 +420,10 @@ func (m *Manager) Rules(r *rulespb.RulesRequest, s rulespb.Rules_RulesServer) er
 	enrichRulesWithExtLabels(pgs, m.extLset)
 
 	for _, pg := range pgs {
-		if err := s.Send(&rulespb.RulesResponse{Result: &rulespb.RulesResponse_Group{Group: pg}}); err != nil {
+		tracing.DoInSpan(s.Context(), "send_rule_group_response", func(_ context.Context) {
+			err = s.Send(&rulespb.RulesResponse{Result: &rulespb.RulesResponse_Group{Group: pg}})
+		})
+		if err != nil {
 			return err
 		}
 	}

pkg/rules/proxy.go

@@ -12,6 +12,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/thanos-io/thanos/pkg/rules/rulespb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
+	"github.com/thanos-io/thanos/pkg/tracing"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -39,10 +40,14 @@ func NewProxy(logger log.Logger, rules func() []rulespb.RulesClient) *Proxy {
 }
 
 func (s *Proxy) Rules(req *rulespb.RulesRequest, srv rulespb.Rules_RulesServer) error {
+	span, ctx := tracing.StartSpan(srv.Context(), "proxy_rules")
+	defer span.Finish()
+
 	var (
-		g, gctx  = errgroup.WithContext(srv.Context())
+		g, gctx  = errgroup.WithContext(ctx)
 		respChan = make(chan *rulespb.RuleGroup, 10)
 		groups   []*rulespb.RuleGroup
+		err      error
 	)
 
 	for _, rulesClient := range s.rules() {
@@ -70,7 +75,10 @@ func (s *Proxy) Rules(req *rulespb.RulesRequest, srv rulespb.Rules_RulesServer)
 	}
 
 	for _, g := range groups {
-		if err := srv.Send(rulespb.NewRuleGroupRulesResponse(g)); err != nil {
+		tracing.DoInSpan(srv.Context(), "send_rules_response", func(_ context.Context) {
+			err = srv.Send(rulespb.NewRuleGroupRulesResponse(g))
+		})
+		if err != nil {
 			return status.Error(codes.Unknown, errors.Wrap(err, "send rules response").Error())
 		}
 	}
@@ -86,7 +94,15 @@ type rulesStream struct {
 }
 
 func (stream *rulesStream) receive(ctx context.Context) error {
-	rules, err := stream.client.Rules(ctx, stream.request)
+	var (
+		err   error
+		rules rulespb.Rules_RulesClient
+	)
+
+	tracing.DoInSpan(ctx, "receive_stream_request", func(ctx context.Context) {
+		rules, err = stream.client.Rules(ctx, stream.request)
+	})
+
 	if err != nil {
 		err = errors.Wrapf(err, "fetching rules from rules client %v", stream.client)
 

pkg/rules/rules.go

@@ -11,6 +11,7 @@ import (
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/thanos-io/thanos/pkg/rules/rulespb"
+	"github.com/thanos-io/thanos/pkg/tracing"
 )
 
 var _ UnaryClient = &GRPCClient{}
@@ -46,6 +47,9 @@ func NewGRPCClientWithDedup(rs rulespb.RulesServer, replicaLabels []string) *GRP
 }
 
 func (rr *GRPCClient) Rules(ctx context.Context, req *rulespb.RulesRequest) (*rulespb.RuleGroups, storage.Warnings, error) {
+	span, ctx := tracing.StartSpan(ctx, "rules_request")
+	defer span.Finish()
+
 	resp := &rulesServer{ctx: ctx}
 
 	if err := rr.proxy.Rules(req, resp); err != nil {