MacroPower/wakatime_exporter

/*
Copyright 2015 The Prometheus Authors
Modifications Copyright 2020 Jacob Colvin (MacroPower)
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package collector

import (
    "errors"
    "fmt"
    "net/url"
    "regexp"
    "strings"
    "sync"
    "time"

    "github.com/go-kit/kit/log"
    "github.com/go-kit/kit/log/level"
    "github.com/prometheus/client_golang/prometheus"
    kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
    scrapeDurationDesc = prometheus.NewDesc(
        prometheus.BuildFQName(namespace, "scrape", "collector_duration_seconds"),
        "wakatime_exporter: Duration of a collector scrape.",
        []string{"collector"},
        nil,
    )
    scrapeSuccessDesc = prometheus.NewDesc(
        prometheus.BuildFQName(namespace, "scrape", "collector_success"),
        "wakatime_exporter: Whether a collector succeeded.",
        []string{"collector"},
        nil,
    )
)

const (
    defaultEnabled  = true
    defaultDisabled = false
)

var (
    factories        = make(map[string]func(in CommonInputs, logger log.Logger) (Collector, error))
    collectorState   = make(map[string]*bool)
    forcedCollectors = map[string]bool{} // collectors which have been explicitly enabled or disabled
)

// CommonInputs are the inputs needed to implement any Collector
type CommonInputs struct {
    BaseURI   url.URL
    URI       url.URL
    Token     string
    SSLVerify bool
    Timeout   time.Duration
}

func registerCollector(collector string, isDefaultEnabled bool, factory func(in CommonInputs, logger log.Logger) (Collector, error)) {
    var helpDefaultState string
    if isDefaultEnabled {
        helpDefaultState = "enabled"
    } else {
        helpDefaultState = "disabled"
    }

    flagName := fmt.Sprintf("collector.%s", collector)
    flagHelp := fmt.Sprintf("Enable the %s collector (default: %s).", collector, helpDefaultState)
    defaultValue := fmt.Sprintf("%v", isDefaultEnabled)

    reg, _ := regexp.Compile("[^a-zA-Z0-9]+")
    envar := "WAKA_COLLECTOR_" + strings.ToUpper(reg.ReplaceAllString(collector, ""))

    flag := kingpin.Flag(flagName, flagHelp).Default(defaultValue).Envar(envar).Action(collectorFlagAction(collector)).Bool()
    collectorState[collector] = flag

    factories[collector] = factory
}

// WakaCollector implements the prometheus.Collector interface.
type WakaCollector struct {
    Collectors map[string]Collector
    logger     log.Logger
}

// DisableDefaultCollectors sets the collector state to false for all collectors which
// have not been explicitly enabled on the command line.
func DisableDefaultCollectors() {
    for c := range collectorState {
        if _, ok := forcedCollectors[c]; !ok {
            *collectorState[c] = false
        }
    }
}

// collectorFlagAction generates a new action function for the given collector
// to track whether it has been explicitly enabled or disabled from the command line.
// A new action function is needed for each collector flag because the ParseContext
// does not contain information about which flag called the action.
// See: https://github.com/alecthomas/kingpin/issues/294
func collectorFlagAction(collector string) func(ctx *kingpin.ParseContext) error {
    return func(ctx *kingpin.ParseContext) error {
        forcedCollectors[collector] = true
        return nil
    }
}

// NewWakaCollector creates a new Collector.
func NewWakaCollector(in CommonInputs, logger log.Logger, filters ...string) (*WakaCollector, error) {
    f := make(map[string]bool)
    for _, filter := range filters {
        enabled, exist := collectorState[filter]
        if !exist {
            return nil, fmt.Errorf("missing collector: %s", filter)
        }
        if !*enabled {
            return nil, fmt.Errorf("disabled collector: %s", filter)
        }
        f[filter] = true
    }
    collectors := make(map[string]Collector)
    for key, enabled := range collectorState {
        if *enabled {
            collector, err := factories[key](in, log.With(logger, "collector", key))
            if err != nil {
                return nil, err
            }
            if len(f) == 0 || f[key] {
                collectors[key] = collector
            }
        }
    }
    return &WakaCollector{Collectors: collectors, logger: logger}, nil
}

// Describe implements the prometheus.Collector interface.
func (n WakaCollector) Describe(ch chan<- *prometheus.Desc) {
    ch <- scrapeDurationDesc
    ch <- scrapeSuccessDesc
}

// Collect implements the prometheus.Collector interface.
func (n WakaCollector) Collect(ch chan<- prometheus.Metric) {
    wg := sync.WaitGroup{}
    wg.Add(len(n.Collectors))
    for name, c := range n.Collectors {
        go func(name string, c Collector) {
            execute(name, c, ch, n.logger)
            wg.Done()
        }(name, c)
    }
    wg.Wait()
}

func execute(name string, c Collector, ch chan<- prometheus.Metric, logger log.Logger) {
    begin := time.Now()
    err := c.Update(ch)
    duration := time.Since(begin)
    var success float64

    if err != nil {
        if isNoDataError(err) {
            level.Debug(logger).Log("msg", "collector returned no data", "name", name, "duration_seconds", duration.Seconds(), "err", err)
        } else {
            level.Error(logger).Log("msg", "collector failed", "name", name, "duration_seconds", duration.Seconds(), "err", err)
        }
        success = 0
    } else {
        level.Debug(logger).Log("msg", "collector succeeded", "name", name, "duration_seconds", duration.Seconds())
        success = 1
    }
    ch <- prometheus.MustNewConstMetric(scrapeDurationDesc, prometheus.GaugeValue, duration.Seconds(), name)
    ch <- prometheus.MustNewConstMetric(scrapeSuccessDesc, prometheus.GaugeValue, success, name)
}

// Collector is the interface a collector has to implement.
type Collector interface {
    // Get new metrics and expose them via prometheus registry.
    Update(ch chan<- prometheus.Metric) error
}

type typedDesc struct {
    desc      *prometheus.Desc
    valueType prometheus.ValueType
}

func (d *typedDesc) mustNewConstMetric(value float64, labels ...string) prometheus.Metric {
    return prometheus.MustNewConstMetric(d.desc, d.valueType, value, labels...)
}

// ErrNoData indicates the collector found no data to collect, but had no other error.
var ErrNoData = errors.New("collector returned no data")

func isNoDataError(err error) bool {
    return err == ErrNoData
}