http.go

Documentation: edge-infra.dev/pkg/lib/promassert

     1  package promassert
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"strings"
     9  
    10  	"github.com/prometheus/client_golang/prometheus/promhttp"
    11  )
    12  
    13  // A barebones http.ResponseWriter implementation to use in the handler's ServeHTTP method.
    14  type responseWriter struct {
    15  	// Embed a bytes.Buffer so that the responseWriter type can be written to and read from.
    16  	// The Write method is needed to implement http.ResponseWriter.
    17  	bytes.Buffer
    18  	header http.Header
    19  	Code   int
    20  }
    21  
    22  // implement http.ResponseWriter
    23  func (rw *responseWriter) Header() http.Header {
    24  	if rw.header == nil {
    25  		rw.header = make(http.Header)
    26  	}
    27  	return rw.header
    28  }
    29  
    30  // implement http.ResponseWriter
    31  func (rw *responseWriter) WriteHeader(statusCode int) {
    32  	rw.Code = statusCode
    33  }
    34  
    35  // Because prometheus is globally oriented, the handler will work in multiple places.
    36  // This handler does not interfere with whatever handler you use in your program.
    37  // We call the handlers ServeHTTP method directly to keep everything in memory.
    38  var promhandler = promhttp.Handler()
    39  
    40  // ScrapePrometheusMetrics allows consumers to provided a custom metrics source. For instance, you may want to scrape from a live HTTP endpoint.
    41  var ScrapePrometheusMetrics = ScrapePromhttpHandler
    42  
    43  // ScrapePromhttpHandler is the default scrape method. It bypasses opening an HTTP port by scraping prometheus data directly from memory.
    44  //
    45  // The ability to scrape metrics in-memory is achieved by sending requests directly to the global promhttp.Handler interface.
    46  func ScrapePromhttpHandler() (io.Reader, error) {
    47  	req, err := http.NewRequest("GET", "http://localhost/", strings.NewReader("")) // the addr is a placeholder and not used.
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  	var rw responseWriter
    52  	promhandler.ServeHTTP(&rw, req)
    53  	if rw.Code != http.StatusOK {
    54  		// Shouldn't ever happen, but need to panic if it does.
    55  		return nil, fmt.Errorf("prometheus metrics HTTP handler did not return 200 OK, got HTTP status code %d", rw.Code)
    56  	}
    57  	return &rw, nil
    58  }
    59  
    60  // ScrapeHTTP returns a func that scrapes the provided HTTP addr each time it is called.
    61  func ScrapeHTTP(addr string) func() (io.Reader, error) {
    62  	return func() (io.Reader, error) {
    63  		resp, err := http.Get(addr) //nolint:gosec // this url is trusted
    64  		if err != nil {
    65  			return nil, err
    66  		}
    67  		defer resp.Body.Close()
    68  		var buf bytes.Buffer
    69  		if _, err := io.Copy(&buf, resp.Body); err != nil {
    70  			return nil, err
    71  		}
    72  		return &buf, nil
    73  	}
    74  }
    75
View as plain text