github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/remotewrite/client.go (about)

     1  package remotewrite
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"fmt"
     8  	"io/ioutil"
     9  	"net/http"
    10  	"strconv"
    11  	"time"
    12  
    13  	"github.com/prometheus/client_golang/prometheus"
    14  	"github.com/sirupsen/logrus"
    15  
    16  	"github.com/pyroscope-io/pyroscope/pkg/config"
    17  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    18  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    19  )
    20  
    21  var (
    22  	ErrConvertPutInputToRequest = errors.New("failed to convert putInput into a http.Request")
    23  	ErrMakingRequest            = errors.New("failed to make request")
    24  	ErrNotOkResponse            = errors.New("response not ok")
    25  )
    26  
    27  type Client struct {
    28  	url     string
    29  	log     logrus.FieldLogger
    30  	config  config.RemoteWriteTarget
    31  	client  *http.Client
    32  	metrics *clientMetrics
    33  }
    34  
    35  func NewClient(logger logrus.FieldLogger, reg prometheus.Registerer, targetName string, cfg config.RemoteWriteTarget) *Client {
    36  	// setup defaults
    37  	if cfg.Timeout == 0 {
    38  		cfg.Timeout = time.Second * 10
    39  	}
    40  
    41  	client := &http.Client{
    42  		Timeout: cfg.Timeout,
    43  		Transport: &http.Transport{
    44  			MaxConnsPerHost:     numWorkers(),
    45  			MaxIdleConns:        numWorkers(),
    46  			MaxIdleConnsPerHost: numWorkers(),
    47  		},
    48  	}
    49  
    50  	metrics := newClientMetrics(reg, targetName, cfg.Address)
    51  	metrics.mustRegister()
    52  
    53  	return &Client{
    54  		url:     cfg.Address + "/ingest",
    55  		log:     logger,
    56  		config:  cfg,
    57  		client:  client,
    58  		metrics: metrics,
    59  	}
    60  }
    61  
    62  func (r *Client) Ingest(ctx context.Context, in *ingestion.IngestInput) error {
    63  	req, err := r.ingestInputToRequest(in)
    64  	if err != nil {
    65  		return fmt.Errorf("%w: %v", ErrConvertPutInputToRequest, err)
    66  	}
    67  
    68  	r.enhanceWithAuth(req)
    69  
    70  	req = req.WithContext(ctx)
    71  
    72  	b, err := in.Profile.Bytes()
    73  	if err == nil {
    74  		// TODO(petethepig): we might want to improve accuracy of this metric at some point
    75  		//   see comment here: https://github.com/pyroscope-io/pyroscope/pull/1147#discussion_r894975126
    76  		r.metrics.sentBytes.Add(float64(len(b)))
    77  	}
    78  
    79  	start := time.Now()
    80  	res, err := r.client.Do(req)
    81  	if res != nil {
    82  		// sometimes both res and err are non-nil
    83  		// therefore we must always close the body first
    84  		defer res.Body.Close()
    85  	}
    86  	if err != nil {
    87  		return fmt.Errorf("%w: %v", ErrMakingRequest, err)
    88  	}
    89  
    90  	duration := time.Since(start)
    91  	r.metrics.responseTime.With(prometheus.Labels{
    92  		"code": strconv.FormatInt(int64(res.StatusCode), 10),
    93  	}).Observe(duration.Seconds())
    94  
    95  	if !(res.StatusCode >= 200 && res.StatusCode < 300) {
    96  		// read all the response body
    97  		respBody, _ := ioutil.ReadAll(res.Body)
    98  		return fmt.Errorf("%w: %v", ErrNotOkResponse, fmt.Errorf("status code: '%d'. body: '%s'. url: '%s'", res.StatusCode, respBody, req.URL.Redacted()))
    99  	}
   100  
   101  	return nil
   102  }
   103  
   104  func (r *Client) ingestInputToRequest(in *ingestion.IngestInput) (*http.Request, error) {
   105  	b, err := in.Profile.Bytes()
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  
   110  	req, err := http.NewRequest("POST", r.url, bytes.NewReader(b))
   111  	for k, v := range r.config.Headers {
   112  		req.Header.Set(k, v)
   113  	}
   114  	if err != nil {
   115  		return nil, err
   116  	}
   117  
   118  	params := req.URL.Query()
   119  	if in.Format != "" {
   120  		params.Set("format", string(in.Format))
   121  	}
   122  
   123  	params.Set("name", r.getQuery(in.Metadata.Key))
   124  	params.Set("from", strconv.FormatInt(in.Metadata.StartTime.Unix(), 10))
   125  	params.Set("until", strconv.FormatInt(in.Metadata.EndTime.Unix(), 10))
   126  	params.Set("sampleRate", strconv.FormatUint(uint64(in.Metadata.SampleRate), 10))
   127  	params.Set("spyName", in.Metadata.SpyName)
   128  	params.Set("units", in.Metadata.Units.String())
   129  	params.Set("aggregationType", in.Metadata.AggregationType.String())
   130  	req.URL.RawQuery = params.Encode()
   131  
   132  	req.Header.Set("Content-Type", in.Profile.ContentType())
   133  
   134  	return req, nil
   135  }
   136  
   137  func (r *Client) getQuery(key *segment.Key) string {
   138  	k := key.Clone()
   139  
   140  	labels := k.Labels()
   141  	for tag, value := range r.config.Tags {
   142  		labels[tag] = value
   143  	}
   144  
   145  	return k.Normalized()
   146  }
   147  
   148  func (r *Client) enhanceWithAuth(req *http.Request) {
   149  	token := r.config.AuthToken
   150  
   151  	if token != "" {
   152  		req.Header.Set("Authorization", "Bearer "+token)
   153  	}
   154  
   155  	if r.config.BasicAuthUser != "" && r.config.BasicAuthPassword != "" {
   156  		req.SetBasicAuth(r.config.BasicAuthUser, r.config.BasicAuthPassword)
   157  	}
   158  	if r.config.TenantID != "" {
   159  		req.Header.Set("X-Scope-OrgID", r.config.TenantID)
   160  	}
   161  }