github.com/grafana/pyroscope@v1.18.0/pkg/test/integration/cluster/component.go (about)

     1  package cluster
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"io"
     8  	"net/http"
     9  	"sync"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  
    13  	"github.com/grafana/pyroscope/pkg/cfg"
    14  	"github.com/grafana/pyroscope/pkg/pyroscope"
    15  )
    16  
    17  type Component struct {
    18  	Target  string
    19  	replica int
    20  	flags   []string
    21  	cfg     pyroscope.Config
    22  	p       *pyroscope.Pyroscope
    23  	reg     *prometheus.Registry
    24  
    25  	httpPort       int
    26  	grpcPort       int
    27  	memberlistPort int
    28  	raftPort       int
    29  }
    30  
    31  func (c *Component) addPorts(ports []int) {
    32  	if len(ports) < 1 {
    33  		return
    34  	}
    35  	c.httpPort = ports[0]
    36  	if len(ports) < 2 {
    37  		return
    38  	}
    39  	c.grpcPort = ports[1]
    40  	if len(ports) < 3 {
    41  		return
    42  	}
    43  	c.memberlistPort = ports[2]
    44  	if len(ports) < 4 {
    45  		return
    46  	}
    47  	c.raftPort = ports[3]
    48  }
    49  
    50  func (comp *Component) querierReadyCheck(ctx context.Context, expectedIngesters, expectedStoreGateways int) (err error) {
    51  	check := comp.checkMetrics().
    52  		addExpectValue(float64(expectedIngesters), "pyroscope_ring_members", "name", "ingester", "state", "ACTIVE").
    53  		addExpectValue(float64(expectedStoreGateways), "pyroscope_ring_members", "name", "store-gateway-client", "state", "ACTIVE")
    54  	return check.run(ctx)
    55  }
    56  
    57  func (comp *Component) distributorReadyCheck(ctx context.Context, expectedIngesters, expectedDistributors, expectedSegmentWriters int) (err error) {
    58  	check := comp.checkMetrics()
    59  	if expectedIngesters > 0 {
    60  		check = check.addExpectValue(float64(expectedIngesters), "pyroscope_ring_members", "name", "ingester", "state", "ACTIVE")
    61  	}
    62  	if expectedSegmentWriters > 0 {
    63  		check = check.addExpectValue(float64(expectedSegmentWriters), "pyroscope_ring_members", "name", "segment-writer", "state", "ACTIVE")
    64  	}
    65  	if expectedDistributors > 0 {
    66  		check = check.addExpectValue(float64(expectedDistributors), "pyroscope_ring_members", "name", "distributor", "state", "ACTIVE")
    67  	}
    68  	return check.run(ctx)
    69  }
    70  
    71  func (comp *Component) httpReadyCheck(ctx context.Context) error {
    72  	req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("http://%s:%d/ready", listenAddr, comp.httpPort), nil)
    73  	if err != nil {
    74  		return err
    75  	}
    76  
    77  	resp, err := http.DefaultClient.Do(req)
    78  	if err != nil {
    79  		return err
    80  	}
    81  
    82  	if resp.StatusCode/100 == 2 {
    83  		return nil
    84  	}
    85  
    86  	body, err := io.ReadAll(resp.Body)
    87  	if err != nil {
    88  		return err
    89  	}
    90  
    91  	return fmt.Errorf("status=%d msg=%s", resp.StatusCode, string(body))
    92  }
    93  
    94  func (comp *Component) Stop() func(context.Context) error {
    95  	return comp.p.Stop()
    96  }
    97  
    98  func (comp *Component) String() string {
    99  	return fmt.Sprintf("[%s] http=%d grpc=%d memberlist=%d raft=%d", comp.nodeName(), comp.httpPort, comp.grpcPort, comp.memberlistPort, comp.raftPort)
   100  }
   101  
   102  func (comp *Component) nodeName() string {
   103  	return fmt.Sprintf("%s-%d", comp.Target, comp.replica)
   104  }
   105  
   106  var lockRegistry sync.Mutex
   107  
   108  func (comp *Component) start(_ context.Context) (*pyroscope.Pyroscope, error) {
   109  	fs := flag.NewFlagSet(comp.nodeName(), flag.PanicOnError)
   110  	if err := cfg.DynamicUnmarshal(&comp.cfg, comp.flags, fs); err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	// Hack to avoid clashing metrics, we should track down the use of globals
   115  	// restore oldReg := prometheus.DefaultRegisterer
   116  	comp.reg = prometheus.NewRegistry()
   117  	lockRegistry.Lock()
   118  	defer lockRegistry.Unlock()
   119  	prometheus.DefaultRegisterer = comp.reg
   120  	prometheus.DefaultGatherer = comp.reg
   121  	comp.cfg.Server.Gatherer = comp.reg
   122  	f, err := pyroscope.New(comp.cfg)
   123  	if err != nil {
   124  		return nil, err
   125  	}
   126  
   127  	return f, nil
   128  }