pcm-coordinator/pkg/tracker/tracker.go

/*

 Copyright (c) [2023] [pcm]
 [pcm-coordinator] is licensed under Mulan PSL v2.
 You can use this software according to the terms and conditions of the Mulan PSL v2.
 You may obtain a copy of Mulan PSL v2 at:
          http://license.coscl.org.cn/MulanPSL2
 THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 EITHER EXPaRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 See the Mulan PSL v2 for more details.

*/

package tracker

import (
	"context"
	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Prometheus struct {
	prometheus Interface
	client     v1.API
}

// NewPrometheus 初始化Prometheus客户端
func NewPrometheus(address string) (Prometheus, error) {
	cfg := api.Config{
		Address: address,
	}

	client, err := api.NewClient(cfg)
	return Prometheus{client: v1.NewAPI(client)}, err
}

func ParseTime(timestamp string) (time.Time, error) {
	// Parse time params
	startInt, err := strconv.ParseInt(timestamp, 10, 64)
	if err != nil {
		return time.Now(), err
	}
	return time.Unix(startInt, 0), nil

}

func (p Prometheus) GetNamedMetricsByTime(metrics []string, start, end string, step time.Duration, o QueryOption) []Metric {
	var res []Metric
	var mtx sync.Mutex
	var wg sync.WaitGroup
	opts := NewQueryOptions()
	o.Apply(opts)

	for _, metric := range metrics {
		wg.Add(1)
		go func(metric string) {
			parsedResp := Metric{MetricName: metric}
			startTimestamp, err := ParseTime(start)
			if err != nil {
				return
			}
			endTimestamp, err := ParseTime(end)
			if err != nil {
				return
			}
			timeRange := v1.Range{
				Start: startTimestamp,
				End:   endTimestamp,
				Step:  step,
			}
			value, _, err := p.client.QueryRange(context.Background(), makeExpr(metric, *opts), timeRange)
			if err != nil {
				parsedResp.Error = err.Error()
			} else {
				parsedResp.MetricData = parseQueryRangeResp(value, genMetricFilter(o))
			}

			mtx.Lock()
			res = append(res, parsedResp)
			mtx.Unlock()

			wg.Done()
		}(metric)
	}

	wg.Wait()

	return res
}

func parseQueryRangeResp(value model.Value, metricFilter func(metric model.Metric) bool) MetricData {
	res := MetricData{MetricType: MetricTypeMatrix}

	data, _ := value.(model.Matrix)

	for _, v := range data {
		if metricFilter != nil && !metricFilter(v.Metric) {
			continue
		}
		mv := MetricValue{
			Metadata: make(map[string]string),
		}

		for k, v := range v.Metric {
			mv.Metadata[string(k)] = string(v)
		}

		for _, k := range v.Values {
			mv.Series = append(mv.Series, Point{float64(k.Timestamp) / 1000, float64(k.Value)})
		}

		res.MetricValues = append(res.MetricValues, mv)
	}

	return res
}

func (p Prometheus) GetNamedMetrics(metrics []string, ts time.Time, o QueryOption) []Metric {
	var res []Metric
	var mtx sync.Mutex
	var wg sync.WaitGroup

	opts := NewQueryOptions()
	o.Apply(opts)

	for _, metric := range metrics {
		wg.Add(1)
		go func(metric string) {
			parsedResp := Metric{MetricName: metric}

			value, _, err := p.client.Query(context.Background(), makeExpr(metric, *opts), ts)
			if err != nil {
				parsedResp.Error = err.Error()
			} else {
				parsedResp.MetricData = parseQueryResp(value, genMetricFilter(o))
			}

			mtx.Lock()
			res = append(res, parsedResp)
			mtx.Unlock()

			wg.Done()
		}(metric)
	}

	wg.Wait()

	return res
}
func parseQueryResp(value model.Value, metricFilter func(metric model.Metric) bool) MetricData {
	res := MetricData{MetricType: MetricTypeVector}

	data, _ := value.(model.Vector)

	for _, v := range data {
		if metricFilter != nil && !metricFilter(v.Metric) {
			continue
		}
		mv := MetricValue{
			Metadata: make(map[string]string),
		}

		for k, v := range v.Metric {
			mv.Metadata[string(k)] = string(v)
		}

		mv.Sample = &Point{float64(v.Timestamp) / 1000, float64(v.Value)}

		res.MetricValues = append(res.MetricValues, mv)
	}

	return res
}

func genMetricFilter(o QueryOption) func(metric model.Metric) bool {
	if o != nil {
		if po, ok := o.(PodOption); ok {
			if po.NamespacedResourcesFilter != "" {
				namespacedPodsMap := make(map[string]struct{})
				for _, s := range strings.Split(po.NamespacedResourcesFilter, "|") {
					namespacedPodsMap[s] = struct{}{}
				}
				return func(metric model.Metric) bool {
					if len(metric) == 0 {
						return false
					}
					_, ok := namespacedPodsMap[string(metric["namespace"])+"/"+string(metric["pod"])]
					return ok
				}
			}
		}
	}
	return func(metric model.Metric) bool {
		return true
	}
}