github.com/thanos-io/thanos@v0.32.5/pkg/metadata/proxy.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package metadata
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/go-kit/log/level"
    12  	"github.com/pkg/errors"
    13  	"golang.org/x/sync/errgroup"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/codes"
    16  	"google.golang.org/grpc/status"
    17  
    18  	"github.com/thanos-io/thanos/pkg/metadata/metadatapb"
    19  	"github.com/thanos-io/thanos/pkg/store/storepb"
    20  	"github.com/thanos-io/thanos/pkg/tracing"
    21  )
    22  
    23  // Proxy implements metadatapb.Metadata gRPC that fanouts requests to given metadatapb.Metadata and deduplication on the way.
    24  type Proxy struct {
    25  	logger   log.Logger
    26  	metadata func() []metadatapb.MetadataClient
    27  }
    28  
    29  func RegisterMetadataServer(metadataSrv metadatapb.MetadataServer) func(*grpc.Server) {
    30  	return func(s *grpc.Server) {
    31  		metadatapb.RegisterMetadataServer(s, metadataSrv)
    32  	}
    33  }
    34  
    35  // NewProxy returns a new metadata.Proxy.
    36  func NewProxy(logger log.Logger, metadata func() []metadatapb.MetadataClient) *Proxy {
    37  	return &Proxy{
    38  		logger:   logger,
    39  		metadata: metadata,
    40  	}
    41  }
    42  
    43  func (s *Proxy) MetricMetadata(req *metadatapb.MetricMetadataRequest, srv metadatapb.Metadata_MetricMetadataServer) error {
    44  	span, ctx := tracing.StartSpan(srv.Context(), "proxy_metadata")
    45  	defer span.Finish()
    46  
    47  	var (
    48  		g, gctx  = errgroup.WithContext(ctx)
    49  		respChan = make(chan *metadatapb.MetricMetadata, 10)
    50  		metas    []*metadatapb.MetricMetadata
    51  		err      error
    52  	)
    53  
    54  	for _, metadataClient := range s.metadata() {
    55  		rs := &metricMetadataStream{
    56  			client:  metadataClient,
    57  			request: req,
    58  			channel: respChan,
    59  			server:  srv,
    60  		}
    61  		g.Go(func() error { return rs.receive(gctx) })
    62  	}
    63  
    64  	go func() {
    65  		_ = g.Wait()
    66  		close(respChan)
    67  	}()
    68  
    69  	for resp := range respChan {
    70  		metas = append(metas, resp)
    71  	}
    72  
    73  	if err := g.Wait(); err != nil {
    74  		level.Error(s.logger).Log("err", err)
    75  		return err
    76  	}
    77  
    78  	for _, t := range metas {
    79  		tracing.DoInSpan(srv.Context(), "send_metadata_response", func(_ context.Context) {
    80  			err = srv.Send(metadatapb.NewMetricMetadataResponse(t))
    81  		})
    82  		if err != nil {
    83  			return status.Error(codes.Unknown, errors.Wrap(err, "send metric metadata response").Error())
    84  		}
    85  	}
    86  
    87  	return nil
    88  }
    89  
    90  type metricMetadataStream struct {
    91  	client  metadatapb.MetadataClient
    92  	request *metadatapb.MetricMetadataRequest
    93  	channel chan<- *metadatapb.MetricMetadata
    94  	server  metadatapb.Metadata_MetricMetadataServer
    95  }
    96  
    97  func (stream *metricMetadataStream) receive(ctx context.Context) error {
    98  	var (
    99  		err         error
   100  		metadataCli metadatapb.Metadata_MetricMetadataClient
   101  	)
   102  
   103  	tracing.DoInSpan(ctx, "receive_metadata_stream_request", func(ctx context.Context) {
   104  		metadataCli, err = stream.client.MetricMetadata(ctx, stream.request)
   105  	})
   106  
   107  	if err != nil {
   108  		err = errors.Wrapf(err, "fetching metric metadata from metadata client %v", stream.client)
   109  
   110  		if stream.request.PartialResponseStrategy == storepb.PartialResponseStrategy_ABORT {
   111  			return err
   112  		}
   113  
   114  		if serr := stream.server.Send(metadatapb.NewWarningMetadataResponse(err)); serr != nil {
   115  			return serr
   116  		}
   117  		// Not an error if response strategy is warning.
   118  		return nil
   119  	}
   120  
   121  	for {
   122  		resp, err := metadataCli.Recv()
   123  		if err == io.EOF {
   124  			return nil
   125  		}
   126  
   127  		if err != nil {
   128  			err = errors.Wrapf(err, "receiving metric metadata from metadata client %v", stream.client)
   129  
   130  			if stream.request.PartialResponseStrategy == storepb.PartialResponseStrategy_ABORT {
   131  				return err
   132  			}
   133  
   134  			if err := stream.server.Send(metadatapb.NewWarningMetadataResponse(err)); err != nil {
   135  				return errors.Wrapf(err, "sending metadata error to server %v", stream.server)
   136  			}
   137  
   138  			// Not an error if response strategy is warning.
   139  			return nil
   140  		}
   141  
   142  		if w := resp.GetWarning(); w != "" {
   143  			if err := stream.server.Send(metadatapb.NewWarningMetadataResponse(errors.New(w))); err != nil {
   144  				return errors.Wrapf(err, "sending metadata warning to server %v", stream.server)
   145  			}
   146  			continue
   147  		}
   148  
   149  		select {
   150  		case stream.channel <- resp.GetMetadata():
   151  		case <-ctx.Done():
   152  			return ctx.Err()
   153  		}
   154  	}
   155  }