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

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package targets
     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/store/storepb"
    19  	"github.com/thanos-io/thanos/pkg/targets/targetspb"
    20  )
    21  
    22  // Proxy implements targetspb.Targets gRPC that fans out requests to given targetspb.Targets.
    23  type Proxy struct {
    24  	logger  log.Logger
    25  	targets func() []targetspb.TargetsClient
    26  }
    27  
    28  func RegisterTargetsServer(targetsSrv targetspb.TargetsServer) func(*grpc.Server) {
    29  	return func(s *grpc.Server) {
    30  		targetspb.RegisterTargetsServer(s, targetsSrv)
    31  	}
    32  }
    33  
    34  // NewProxy returns new targets.Proxy.
    35  func NewProxy(logger log.Logger, targets func() []targetspb.TargetsClient) *Proxy {
    36  	return &Proxy{
    37  		logger:  logger,
    38  		targets: targets,
    39  	}
    40  }
    41  
    42  func (s *Proxy) Targets(req *targetspb.TargetsRequest, srv targetspb.Targets_TargetsServer) error {
    43  	var (
    44  		g, gctx  = errgroup.WithContext(srv.Context())
    45  		respChan = make(chan *targetspb.TargetDiscovery, 10)
    46  		targets  []*targetspb.TargetDiscovery
    47  	)
    48  
    49  	for _, targetsClient := range s.targets() {
    50  		rs := &targetsStream{
    51  			client:  targetsClient,
    52  			request: req,
    53  			channel: respChan,
    54  			server:  srv,
    55  		}
    56  		g.Go(func() error { return rs.receive(gctx) })
    57  	}
    58  
    59  	go func() {
    60  		_ = g.Wait()
    61  		close(respChan)
    62  	}()
    63  
    64  	for resp := range respChan {
    65  		// TODO: Stream it
    66  		targets = append(targets, resp)
    67  	}
    68  
    69  	if err := g.Wait(); err != nil {
    70  		level.Error(s.logger).Log("err", err)
    71  		return err
    72  	}
    73  
    74  	for _, t := range targets {
    75  		if err := srv.Send(targetspb.NewTargetsResponse(t)); err != nil {
    76  			return status.Error(codes.Unknown, errors.Wrap(err, "send targets response").Error())
    77  		}
    78  	}
    79  
    80  	return nil
    81  }
    82  
    83  type targetsStream struct {
    84  	client  targetspb.TargetsClient
    85  	request *targetspb.TargetsRequest
    86  	channel chan<- *targetspb.TargetDiscovery
    87  	server  targetspb.Targets_TargetsServer
    88  }
    89  
    90  func (stream *targetsStream) receive(ctx context.Context) error {
    91  	targets, err := stream.client.Targets(ctx, stream.request)
    92  	if err != nil {
    93  		err = errors.Wrapf(err, "fetching targets from targets client %v", stream.client)
    94  
    95  		if stream.request.PartialResponseStrategy == storepb.PartialResponseStrategy_ABORT {
    96  			return err
    97  		}
    98  
    99  		if serr := stream.server.Send(targetspb.NewWarningTargetsResponse(err)); serr != nil {
   100  			return serr
   101  		}
   102  		// Not an error if response strategy is warning.
   103  		return nil
   104  	}
   105  
   106  	for {
   107  		target, err := targets.Recv()
   108  		if err == io.EOF {
   109  			return nil
   110  		}
   111  
   112  		if err != nil {
   113  			err = errors.Wrapf(err, "receiving targets from targets client %v", stream.client)
   114  
   115  			if stream.request.PartialResponseStrategy == storepb.PartialResponseStrategy_ABORT {
   116  				return err
   117  			}
   118  
   119  			if err := stream.server.Send(targetspb.NewWarningTargetsResponse(err)); err != nil {
   120  				return errors.Wrapf(err, "sending targets error to server %v", stream.server)
   121  			}
   122  			// Not an error if response strategy is warning.
   123  			return nil
   124  		}
   125  
   126  		if w := target.GetWarning(); w != "" {
   127  			if err := stream.server.Send(targetspb.NewWarningTargetsResponse(errors.New(w))); err != nil {
   128  				return errors.Wrapf(err, "sending targets warning to server %v", stream.server)
   129  			}
   130  			continue
   131  		}
   132  
   133  		select {
   134  		case stream.channel <- target.GetTargets():
   135  		case <-ctx.Done():
   136  			return ctx.Err()
   137  		}
   138  	}
   139  }