From aa7ae5fc65e41321ad5aa7400ac79f3c00b6ed68 Mon Sep 17 00:00:00 2001
From: Will Assis <35489495+gassiss@users.noreply.github.com>
Date: Mon, 28 Jul 2025 12:35:20 -0400
Subject: [PATCH] unified-storage: add tracing to distributor methods (#108791)

* add tracing to distributor methods
---
 .../unified/resource/search_server_distributor.go      | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/pkg/storage/unified/resource/search_server_distributor.go b/pkg/storage/unified/resource/search_server_distributor.go
index 14527f0e6b3..1e1c00e7a7e 100644
--- a/pkg/storage/unified/resource/search_server_distributor.go
+++ b/pkg/storage/unified/resource/search_server_distributor.go
@@ -33,6 +33,7 @@ func ProvideSearchDistributorServer(cfg *setting.Cfg, features featuremgmt.Featu
 		log:        log.New("index-server-distributor"),
 		ring:       ring,
 		clientPool: ringClientPool,
+		tracing:    tracer,
 	}
 
 	healthService, err := ProvideHealthService(distributorServer)
@@ -80,6 +81,7 @@ type distributorServer struct {
 	clientPool *ringclient.Pool
 	ring       *ring.Ring
 	log        log.Logger
+	tracing    trace.Tracer
 }
 
 var (
@@ -92,6 +94,8 @@ var (
 )
 
 func (ds *distributorServer) Search(ctx context.Context, r *resourcepb.ResourceSearchRequest) (*resourcepb.ResourceSearchResponse, error) {
+	ctx, span := ds.tracing.Start(ctx, "distributor.Search")
+	defer span.End()
 	ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Options.Key.Namespace, "Search")
 	if err != nil {
 		return nil, err
@@ -101,6 +105,8 @@ func (ds *distributorServer) Search(ctx context.Context, r *resourcepb.ResourceS
 }
 
 func (ds *distributorServer) GetStats(ctx context.Context, r *resourcepb.ResourceStatsRequest) (*resourcepb.ResourceStatsResponse, error) {
+	ctx, span := ds.tracing.Start(ctx, "distributor.GetStats")
+	defer span.End()
 	ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Namespace, "GetStats")
 	if err != nil {
 		return nil, err
@@ -110,6 +116,8 @@ func (ds *distributorServer) GetStats(ctx context.Context, r *resourcepb.Resourc
 }
 
 func (ds *distributorServer) CountManagedObjects(ctx context.Context, r *resourcepb.CountManagedObjectsRequest) (*resourcepb.CountManagedObjectsResponse, error) {
+	ctx, span := ds.tracing.Start(ctx, "distributor.CountManagedObjects")
+	defer span.End()
 	ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Namespace, "CountManagedObjects")
 	if err != nil {
 		return nil, err
@@ -119,6 +127,8 @@ func (ds *distributorServer) CountManagedObjects(ctx context.Context, r *resourc
 }
 
 func (ds *distributorServer) ListManagedObjects(ctx context.Context, r *resourcepb.ListManagedObjectsRequest) (*resourcepb.ListManagedObjectsResponse, error) {
+	ctx, span := ds.tracing.Start(ctx, "distributor.ListManagedObjects")
+	defer span.End()
 	ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Namespace, "ListManagedObjects")
 	if err != nil {
 		return nil, err