github.com/m-lab/locate@v0.17.6/handler/prometheus.go (about)

     1  package handler
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"log"
     8  	"net/http"
     9  	"time"
    10  
    11  	"github.com/m-lab/go/host"
    12  	"github.com/m-lab/locate/static"
    13  	prom "github.com/prometheus/client_golang/api/prometheus/v1"
    14  	"github.com/prometheus/common/model"
    15  )
    16  
    17  var (
    18  	timeout         = static.PrometheusCheckPeriod
    19  	errCouldNotCast = errors.New("could not cast metric to vector")
    20  
    21  	// End-to-end query parameters.
    22  	// The raw query should be unfiltered.
    23  	e2eQuery = "script_success"
    24  	e2eLabel = model.LabelName("fqdn")
    25  	// The script was successful if the value != 0.
    26  	e2eFunction = func(v float64) bool {
    27  		return v != 0
    28  	}
    29  
    30  	// GMX query parameters.
    31  	gmxQuery = "gmx_machine_maintenance"
    32  	// The raw query should be unfiltered.
    33  	gmxLabel = model.LabelName("machine")
    34  	// The machine is not in maintenance if the value = 0.
    35  	gmxFunction = func(v float64) bool {
    36  		return v == 0
    37  	}
    38  )
    39  
    40  // Prometheus is a handler that collects Prometheus health signals.
    41  func (c *Client) Prometheus(rw http.ResponseWriter, req *http.Request) {
    42  	err := c.updatePrometheus(req.Context(), "")
    43  	if err != nil {
    44  		rw.WriteHeader(http.StatusInternalServerError)
    45  		return
    46  	}
    47  
    48  	rw.WriteHeader(http.StatusOK)
    49  }
    50  
    51  // UpdatePrometheusForMachine updates the Prometheus signals for a single machine hostname.
    52  func (c *Client) UpdatePrometheusForMachine(ctx context.Context, hostname string) error {
    53  	name, err := host.Parse(hostname)
    54  	if err != nil {
    55  		log.Printf("Error parsing hostname %s", hostname)
    56  		return err
    57  	}
    58  
    59  	machine := name.String()
    60  	err = c.updatePrometheus(ctx, fmt.Sprintf("machine=%q", machine))
    61  	if err != nil {
    62  		log.Printf("Error updating Prometheus signals for machine %s", machine)
    63  	}
    64  	return err
    65  }
    66  
    67  func (c *Client) updatePrometheus(ctx context.Context, filter string) error {
    68  	hostnames, err := c.query(ctx, e2eQuery, filter, e2eLabel, e2eFunction)
    69  	if err != nil {
    70  		log.Printf("Error querying Prometheus for %s metric: %v", e2eQuery, err)
    71  		return err
    72  	}
    73  
    74  	machines, err := c.query(ctx, gmxQuery, filter, gmxLabel, gmxFunction)
    75  	if err != nil {
    76  		log.Printf("Error querying Prometheus for %s metric: %v", gmxQuery, err)
    77  		return err
    78  	}
    79  
    80  	err = c.UpdatePrometheus(hostnames, machines)
    81  	if err != nil {
    82  		log.Printf("Error updating internal Prometheus state: %v", err)
    83  		return err
    84  	}
    85  
    86  	return nil
    87  }
    88  
    89  // query performs the provided PromQL query.
    90  func (c *Client) query(ctx context.Context, query, filter string, labelName model.LabelName, f func(v float64) bool) (map[string]bool, error) {
    91  	result, _, err := c.PrometheusClient.Query(ctx, formatQuery(query, filter), time.Now(), prom.WithTimeout(timeout))
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  
    96  	vector, ok := result.(model.Vector)
    97  	if !ok {
    98  		return nil, errCouldNotCast
    99  	}
   100  
   101  	return getMetrics(vector, labelName, f), nil
   102  }
   103  
   104  // getMetrics returns a map of labels to bool values from a Vector, based on the function parameter.
   105  func getMetrics(vector model.Vector, labelName model.LabelName, f func(v float64) bool) map[string]bool {
   106  	metrics := map[string]bool{}
   107  
   108  	for _, elem := range vector {
   109  		label := string(elem.Metric[labelName])
   110  		value := float64(elem.Value)
   111  		metrics[label] = f(value)
   112  	}
   113  
   114  	return metrics
   115  }
   116  
   117  func formatQuery(query, filter string) string {
   118  	if filter != "" {
   119  		return fmt.Sprintf("%s{%s}", query, filter)
   120  	}
   121  	return query
   122  }