github.com/thanos-io/thanos@v0.32.5/pkg/metadata/metadata.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  	"sync"
     9  
    10  	"github.com/pkg/errors"
    11  	"github.com/prometheus/prometheus/storage"
    12  	"github.com/thanos-io/thanos/pkg/metadata/metadatapb"
    13  	"github.com/thanos-io/thanos/pkg/tracing"
    14  )
    15  
    16  var _ UnaryClient = &GRPCClient{}
    17  
    18  // UnaryClient is a gRPC metadatapb.Metadata client which expands streaming metadata API. Useful for consumers that does not
    19  // support streaming.
    20  type UnaryClient interface {
    21  	MetricMetadata(ctx context.Context, req *metadatapb.MetricMetadataRequest) (map[string][]metadatapb.Meta, storage.Warnings, error)
    22  }
    23  
    24  // GRPCClient allows to retrieve metadata from local gRPC streaming server implementation.
    25  // TODO(bwplotka): Switch to native gRPC transparent client->server adapter once available.
    26  type GRPCClient struct {
    27  	proxy metadatapb.MetadataServer
    28  }
    29  
    30  func NewGRPCClient(ts metadatapb.MetadataServer) *GRPCClient {
    31  	return &GRPCClient{
    32  		proxy: ts,
    33  	}
    34  }
    35  
    36  func (rr *GRPCClient) MetricMetadata(ctx context.Context, req *metadatapb.MetricMetadataRequest) (map[string][]metadatapb.Meta, storage.Warnings, error) {
    37  	span, ctx := tracing.StartSpan(ctx, "metadata_grpc_request")
    38  	defer span.Finish()
    39  
    40  	srv := &metadataServer{ctx: ctx, metric: req.Metric, limit: int(req.Limit)}
    41  
    42  	if req.Limit >= 0 {
    43  		if req.Metric != "" {
    44  			srv.metadataMap = make(map[string][]metadatapb.Meta, 1)
    45  		} else if req.Limit <= 100 {
    46  			srv.metadataMap = make(map[string][]metadatapb.Meta, req.Limit)
    47  		} else {
    48  			srv.metadataMap = make(map[string][]metadatapb.Meta)
    49  		}
    50  	} else {
    51  		srv.metadataMap = make(map[string][]metadatapb.Meta)
    52  	}
    53  
    54  	if err := rr.proxy.MetricMetadata(req, srv); err != nil {
    55  		return nil, nil, errors.Wrap(err, "proxy MetricMetadata")
    56  	}
    57  
    58  	return srv.metadataMap, srv.warnings, nil
    59  }
    60  
    61  type metadataServer struct {
    62  	// This field just exist to pseudo-implement the unused methods of the interface.
    63  	metadatapb.Metadata_MetricMetadataServer
    64  	ctx context.Context
    65  
    66  	metric string
    67  	limit  int
    68  
    69  	warnings    []error
    70  	metadataMap map[string][]metadatapb.Meta
    71  	mu          sync.Mutex
    72  }
    73  
    74  func (srv *metadataServer) Send(res *metadatapb.MetricMetadataResponse) error {
    75  	if res.GetWarning() != "" {
    76  		srv.mu.Lock()
    77  		defer srv.mu.Unlock()
    78  		srv.warnings = append(srv.warnings, errors.New(res.GetWarning()))
    79  		return nil
    80  	}
    81  
    82  	if res.GetMetadata() == nil {
    83  		return errors.New("no metadata")
    84  	}
    85  
    86  	// If limit is set to 0, we don't need to add anything.
    87  	if srv.limit == 0 {
    88  		return nil
    89  	}
    90  
    91  	srv.mu.Lock()
    92  	defer srv.mu.Unlock()
    93  	for k, v := range res.GetMetadata().Metadata {
    94  		if metadata, ok := srv.metadataMap[k]; !ok {
    95  			// If limit is set and it is positive, we limit the size of the map.
    96  			if srv.limit < 0 || srv.limit > 0 && len(srv.metadataMap) < srv.limit {
    97  				srv.metadataMap[k] = v.Metas
    98  			}
    99  		} else {
   100  			// There shouldn't be many metadata for one single metric.
   101  		Outer:
   102  			for _, meta := range v.Metas {
   103  				for _, m := range metadata {
   104  					if meta == m {
   105  						continue Outer
   106  					}
   107  				}
   108  				srv.metadataMap[k] = append(srv.metadataMap[k], meta)
   109  			}
   110  		}
   111  	}
   112  
   113  	return nil
   114  }
   115  
   116  func (srv *metadataServer) Context() context.Context {
   117  	return srv.ctx
   118  }