github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/prometheus.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  
     3  package runner
     4  
     5  import (
     6  	"io/ioutil"
     7  	"net/http"
     8  	"strconv"
     9  	"strings"
    10  
    11  	"github.com/go-stack/stack"
    12  	"github.com/jjeffery/kv" // MIT License
    13  
    14  	dto "github.com/prometheus/client_model/go"
    15  	"github.com/prometheus/common/expfmt"
    16  )
    17  
    18  // This file contains the implementation of a prometheus test
    19  // probe that will check the servers metrics resource for metrics
    20  // data that test cases need to validate expected behavior within
    21  // the server logic
    22  
    23  type prometheusClient struct {
    24  	url string
    25  }
    26  
    27  // NewPrometheusClient will instantiate the structure used to communicate with a
    28  // remote prometheus endpoint
    29  //
    30  func NewPrometheusClient(url string) (cli *prometheusClient) {
    31  	return &prometheusClient{
    32  		url: url,
    33  	}
    34  }
    35  
    36  // Fetch will return the family of metrics from prometheus that have the supplied prefix.
    37  //
    38  func (p *prometheusClient) Fetch(prefix string) (metrics map[string]*dto.MetricFamily, err kv.Error) {
    39  	metrics = map[string]*dto.MetricFamily{}
    40  
    41  	resp, errGo := http.Get(p.url)
    42  	if errGo != nil {
    43  		return metrics, kv.Wrap(errGo).With("URL", p.url).With("stack", stack.Trace().TrimRuntime())
    44  	}
    45  	defer resp.Body.Close()
    46  
    47  	parser := expfmt.TextParser{}
    48  	metricFamilies, errGo := parser.TextToMetricFamilies(resp.Body)
    49  	if errGo != nil {
    50  		return metrics, kv.Wrap(errGo).With("URL", p.url).With("stack", stack.Trace().TrimRuntime())
    51  	}
    52  	for k, v := range metricFamilies {
    53  		if len(prefix) == 0 || strings.HasPrefix(k, prefix) {
    54  			metrics[k] = v
    55  		}
    56  	}
    57  	return metrics, nil
    58  }
    59  
    60  func (p *prometheusClient) getMetric(prefix string) (items []string, err kv.Error) {
    61  
    62  	resp, errGo := http.Get(p.url)
    63  	if errGo != nil {
    64  		return items, kv.Wrap(errGo).With("URL", p.url).With("stack", stack.Trace().TrimRuntime())
    65  	}
    66  	defer resp.Body.Close()
    67  
    68  	body, errGo := ioutil.ReadAll(resp.Body)
    69  	if errGo != nil {
    70  		return items, kv.Wrap(errGo).With("URL", p.url).With("stack", stack.Trace().TrimRuntime())
    71  	}
    72  
    73  	lines := strings.Split(string(body), "\n")
    74  	for _, line := range lines {
    75  		if len(prefix) == 0 || strings.HasPrefix(line, prefix) {
    76  			items = append(items, line)
    77  		}
    78  	}
    79  	return items, nil
    80  }
    81  
    82  // GetHitsMisses is a convineance method to get cache hits and misses for runner and StudioML
    83  // artifacts.
    84  //
    85  func (p *prometheusClient) GetHitsMisses(hash string) (hits int, misses int, err kv.Error) {
    86  	lines, err := p.getMetric("runner_cache")
    87  	if err != nil {
    88  		return hits, misses, err
    89  	}
    90  	hashData := "hash=\"" + hash + "\""
    91  	for _, line := range lines {
    92  		if strings.Contains(line, hashData) && strings.HasPrefix(line, "runner_cache") {
    93  			values := strings.Split(line, " ")
    94  			switch {
    95  			case strings.HasPrefix(line, "runner_cache_hits{"):
    96  				hits, _ = strconv.Atoi(values[len(values)-1])
    97  			case strings.HasPrefix(line, "runner_cache_misses{"):
    98  				misses, _ = strconv.Atoi(values[len(values)-1])
    99  			}
   100  		}
   101  	}
   102  	return hits, misses, nil
   103  }