grafana/pkg/storage/unified/resource/search_server_distributor.go

package resource

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/grafana/dskit/ring"
	ringclient "github.com/grafana/dskit/ring/client"
	"github.com/grafana/dskit/services"
	userutils "github.com/grafana/dskit/user"
	"github.com/prometheus/client_golang/prometheus"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/metadata"

	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/grpcserver"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
)

func ProvideSearchDistributorServer(cfg *setting.Cfg, features featuremgmt.FeatureToggles, registerer prometheus.Registerer, tracer trace.Tracer, ring *ring.Ring, ringClientPool *ringclient.Pool) (grpcserver.Provider, error) {
	var err error
	grpcHandler, err := grpcserver.ProvideService(cfg, features, nil, tracer, registerer)
	if err != nil {
		return nil, err
	}

	distributorServer := &distributorServer{
		log:        log.New("index-server-distributor"),
		ring:       ring,
		clientPool: ringClientPool,
		tracing:    tracer,
	}

	healthService, err := ProvideHealthService(distributorServer)
	if err != nil {
		return nil, err
	}

	grpcServer := grpcHandler.GetServer()

	resourcepb.RegisterResourceIndexServer(grpcServer, distributorServer)
	resourcepb.RegisterManagedObjectIndexServer(grpcServer, distributorServer)
	grpc_health_v1.RegisterHealthServer(grpcServer, healthService)
	_, err = grpcserver.ProvideReflectionService(cfg, grpcHandler)
	if err != nil {
		return nil, err
	}

	return grpcHandler, nil
}

type RingClient struct {
	Client ResourceClient
	grpc_health_v1.HealthClient
	Conn *grpc.ClientConn
}

func (c *RingClient) Close() error {
	return c.Conn.Close()
}

func (c *RingClient) String() string {
	return c.RemoteAddress()
}

func (c *RingClient) RemoteAddress() string {
	return c.Conn.Target()
}

const RingKey = "search-server-ring"
const RingName = "search_server_ring"
const RingHeartbeatTimeout = time.Minute
const RingNumTokens = 128

type distributorServer struct {
	clientPool *ringclient.Pool
	ring       *ring.Ring
	log        log.Logger
	tracing    trace.Tracer
}

var (
	// operation used by the distributor to select only ACTIVE instances to handle search-related requests
	searchRingRead = ring.NewOp([]ring.InstanceState{ring.ACTIVE}, func(s ring.InstanceState) bool {
		return s != ring.ACTIVE
	})
)

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
	}

	return client.Search(ctx, r)
}

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
	}

	return client.GetStats(ctx, r)
}

func (ds *distributorServer) RebuildIndexes(ctx context.Context, r *resourcepb.RebuildIndexesRequest) (*resourcepb.RebuildIndexesResponse, error) {
	ctx, span := ds.tracing.Start(ctx, "distributor.RebuildIndexes")
	defer span.End()

	// validate input
	for _, key := range r.Keys {
		if r.Namespace != key.Namespace {
			return &resourcepb.RebuildIndexesResponse{
				Error: NewBadRequestError("key namespace does not match request namespace"),
			}, nil
		}
	}

	// distribute the request to all search pods to minimize risk of stale index
	// it will not rebuild on those which don't have the index open
	rs, err := ds.ring.GetAllHealthy(searchRingRead)
	if err != nil {
		return nil, fmt.Errorf("failed to get all healthy instances from the ring")
	}

	err = grpc.SetHeader(ctx, metadata.Pairs("proxied-instance-id", "all"))
	if err != nil {
		ds.log.Debug("error setting grpc header", "err", err)
	}

	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = make(metadata.MD)
	}
	rCtx := userutils.InjectOrgID(metadata.NewOutgoingContext(ctx, md), r.Namespace)

	var wg sync.WaitGroup
	var totalRebuildCount atomic.Int64
	detailsCh := make(chan string, len(rs.Instances))
	errorCh := make(chan error, len(rs.Instances))

	for _, inst := range rs.Instances {
		wg.Add(1)
		go func() {
			defer wg.Done()

			client, err := ds.clientPool.GetClientForInstance(inst)
			if err != nil {
				errorCh <- fmt.Errorf("instance %s: failed to get client, %w", inst.Id, err)
				return
			}

			rsp, err := client.(*RingClient).Client.RebuildIndexes(rCtx, r)
			if err != nil {
				errorCh <- fmt.Errorf("instance %s: failed to distribute rebuild index request, %w", inst.Id, err)
				return
			}

			if rsp.Error != nil {
				errorCh <- fmt.Errorf("instance %s: rebuild index request returned the error %s", inst.Id, rsp.Error.Message)
				return
			}

			if rsp.Details != "" {
				detailsCh <- fmt.Sprintf("{instance: %s, details: %s}", inst.Id, rsp.Details)
			}

			totalRebuildCount.Add(rsp.RebuildCount)
		}()
	}

	wg.Wait()
	close(errorCh)
	close(detailsCh)

	errs := make([]error, 0, len(errorCh))
	for err := range errorCh {
		ds.log.Error("rebuild indexes call failed with %w", err)
		errs = append(errs, err)
	}

	var details string
	for d := range detailsCh {
		if len(details) > 0 {
			details += ", "
		}
		details += d
	}

	response := &resourcepb.RebuildIndexesResponse{
		RebuildCount: totalRebuildCount.Load(),
		Details:      details,
	}
	if len(errs) > 0 {
		response.Error = AsErrorResult(errors.Join(errs...))
	}
	return response, nil
}

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
	}

	return client.CountManagedObjects(ctx, r)
}

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
	}

	return client.ListManagedObjects(ctx, r)
}

func (ds *distributorServer) getClientToDistributeRequest(ctx context.Context, namespace string, methodName string) (context.Context, ResourceClient, error) {
	ringHasher := fnv.New32a()
	_, err := ringHasher.Write([]byte(namespace))
	if err != nil {
		ds.log.Debug("error hashing namespace", "err", err, "namespace", namespace)
		return ctx, nil, err
	}

	rs, err := ds.ring.GetWithOptions(ringHasher.Sum32(), searchRingRead, ring.WithReplicationFactor(ds.ring.ReplicationFactor()))
	if err != nil {
		ds.log.Debug("error getting replication set from ring", "err", err, "namespace", namespace)
		return ctx, nil, err
	}

	// Randomly select an instance for primitive load balancing
	inst := rs.Instances[rand.Intn(len(rs.Instances))]
	client, err := ds.clientPool.GetClientForInstance(inst)
	if err != nil {
		ds.log.Debug("error getting instance client from pool", "err", err, "namespace", namespace, "searchApiInstanceId", inst.Id)
		return ctx, nil, err
	}

	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = make(metadata.MD)
	}

	err = grpc.SetHeader(ctx, metadata.Pairs("proxied-instance-id", inst.Id))
	if err != nil {
		ds.log.Debug("error setting grpc header", "err", err)
	}

	return userutils.InjectOrgID(metadata.NewOutgoingContext(ctx, md), namespace), client.(*RingClient).Client, nil
}

func (ds *distributorServer) IsHealthy(ctx context.Context, r *resourcepb.HealthCheckRequest) (*resourcepb.HealthCheckResponse, error) {
	if ds.ring.State() == services.Running {
		return &resourcepb.HealthCheckResponse{Status: resourcepb.HealthCheckResponse_SERVING}, nil
	}

	return &resourcepb.HealthCheckResponse{Status: resourcepb.HealthCheckResponse_NOT_SERVING}, nil
}