github.com/netdata/go.d.plugin@v0.58.1/pkg/prometheus/client.go (about)

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package prometheus
     4  
     5  import (
     6  	"bufio"
     7  	"bytes"
     8  	"compress/gzip"
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"net/url"
    13  	"os"
    14  	"path/filepath"
    15  
    16  	"github.com/netdata/go.d.plugin/pkg/prometheus/selector"
    17  	"github.com/netdata/go.d.plugin/pkg/web"
    18  )
    19  
    20  type (
    21  	// Prometheus is a helper for scrape and parse prometheus format metrics.
    22  	Prometheus interface {
    23  		// ScrapeSeries and parse prometheus format metrics
    24  		ScrapeSeries() (Series, error)
    25  		Scrape() (MetricFamilies, error)
    26  		HTTPClient() *http.Client
    27  	}
    28  
    29  	prometheus struct {
    30  		client   *http.Client
    31  		request  web.Request
    32  		filepath string
    33  
    34  		sr selector.Selector
    35  
    36  		parser promTextParser
    37  
    38  		buf     *bytes.Buffer
    39  		gzipr   *gzip.Reader
    40  		bodyBuf *bufio.Reader
    41  	}
    42  )
    43  
    44  const (
    45  	acceptHeader    = `text/plain;version=0.0.4;q=1,*/*;q=0.1`
    46  	userAgentHeader = `netdata/go.d.plugin`
    47  )
    48  
    49  // New creates a Prometheus instance.
    50  func New(client *http.Client, request web.Request) Prometheus {
    51  	return &prometheus{
    52  		client:  client,
    53  		request: request,
    54  		buf:     bytes.NewBuffer(make([]byte, 0, 16000)),
    55  	}
    56  }
    57  
    58  // NewWithSelector creates a Prometheus instance with the selector.
    59  func NewWithSelector(client *http.Client, request web.Request, sr selector.Selector) Prometheus {
    60  	p := &prometheus{
    61  		client:  client,
    62  		request: request,
    63  		sr:      sr,
    64  		buf:     bytes.NewBuffer(make([]byte, 0, 16000)),
    65  		parser:  promTextParser{sr: sr},
    66  	}
    67  
    68  	if v, err := url.Parse(request.URL); err == nil && v.Scheme == "file" {
    69  		p.filepath = filepath.Join(v.Host, v.Path)
    70  	}
    71  
    72  	return p
    73  }
    74  
    75  func (p *prometheus) HTTPClient() *http.Client {
    76  	return p.client
    77  }
    78  
    79  // ScrapeSeries scrapes metrics, parses and sorts
    80  func (p *prometheus) ScrapeSeries() (Series, error) {
    81  	p.buf.Reset()
    82  
    83  	if err := p.fetch(p.buf); err != nil {
    84  		return nil, err
    85  	}
    86  
    87  	return p.parser.parseToSeries(p.buf.Bytes())
    88  }
    89  
    90  func (p *prometheus) Scrape() (MetricFamilies, error) {
    91  	p.buf.Reset()
    92  
    93  	if err := p.fetch(p.buf); err != nil {
    94  		return nil, err
    95  	}
    96  
    97  	return p.parser.parseToMetricFamilies(p.buf.Bytes())
    98  }
    99  
   100  func (p *prometheus) fetch(w io.Writer) error {
   101  	// TODO: should be a separate text file prom client
   102  	if p.filepath != "" {
   103  		f, err := os.Open(p.filepath)
   104  		if err != nil {
   105  			return err
   106  		}
   107  		defer f.Close()
   108  
   109  		_, err = io.Copy(w, f)
   110  
   111  		return err
   112  	}
   113  
   114  	req, err := web.NewHTTPRequest(p.request)
   115  	if err != nil {
   116  		return err
   117  	}
   118  
   119  	req.Header.Add("Accept", acceptHeader)
   120  	req.Header.Add("Accept-Encoding", "gzip")
   121  	req.Header.Set("User-Agent", userAgentHeader)
   122  
   123  	resp, err := p.client.Do(req)
   124  	if err != nil {
   125  		return err
   126  	}
   127  
   128  	defer func() {
   129  		_, _ = io.Copy(io.Discard, resp.Body)
   130  		_ = resp.Body.Close()
   131  	}()
   132  
   133  	if resp.StatusCode != http.StatusOK {
   134  		return fmt.Errorf("server '%s' returned HTTP status code %d (%s)", req.URL, resp.StatusCode, resp.Status)
   135  	}
   136  
   137  	if resp.Header.Get("Content-Encoding") != "gzip" {
   138  		_, err = io.Copy(w, resp.Body)
   139  		return err
   140  	}
   141  
   142  	if p.gzipr == nil {
   143  		p.bodyBuf = bufio.NewReader(resp.Body)
   144  		p.gzipr, err = gzip.NewReader(p.bodyBuf)
   145  		if err != nil {
   146  			return err
   147  		}
   148  	} else {
   149  		p.bodyBuf.Reset(resp.Body)
   150  		_ = p.gzipr.Reset(p.bodyBuf)
   151  	}
   152  
   153  	_, err = io.Copy(w, p.gzipr)
   154  	_ = p.gzipr.Close()
   155  
   156  	return err
   157  }