<commit_before>package prometheus_output

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	dto "github.com/prometheus/client_model/go"
	"google.golang.org/protobuf/proto"
)

const (
	defaultListen     = ":9804"
	defaultPath       = "/metrics"
	defaultExpiration = time.Minute
	defaultMetricHelp = "gNMIc generated metric"
)

type labelPair struct {
	Name  string
	Value string
}
type promMetric struct {
	name   string
	labels []*labelPair
	time   time.Time
	value  float64
}

func init() {
	outputs.Register("prometheus", func() outputs.Output {
		return &PrometheusOutput{
			Cfg:       &Config{},
			eventChan: make(chan *collector.EventMsg),
			wg:        new(sync.WaitGroup),
			entries:   make(map[uint64]*promMetric),
		}
	})
}

type PrometheusOutput struct {
	Cfg       *Config
	metrics   []prometheus.Collector
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *collector.EventMsg

	wg     *sync.WaitGroup
	server *http.Server
	sync.Mutex
	entries map[uint64]*promMetric

	replacer *strings.Replacer
}
type Config struct {
	Listen     string        `mapstructure:"listen,omitempty"`
	Path       string        `mapstructure:"path,omitempty"`
	Expiration time.Duration `mapstructure:"expiration,omitempty"`
	Debug      bool          `mapstructure:"debug,omitempty"`
}

func (p *PrometheusOutput) String() string {
	b, err := json.Marshal(p)
	if err != nil {
		return ""
	}
	return string(b)
}
func (p *PrometheusOutput) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	ctx, p.cancelFn = context.WithCancel(ctx)
	decoder, err := mapstructure.NewDecoder(
		&mapstructure.DecoderConfig{
			DecodeHook: mapstructure.StringToTimeDurationHookFunc(),
			Result:     p.Cfg,
		},
	)
	if err != nil {
		return err
	}
	err = decoder.Decode(cfg)
	if err != nil {
		return err
	}
	if p.Cfg.Listen == "" {
		p.Cfg.Listen = defaultListen
	}
	if p.Cfg.Path == "" {
		p.Cfg.Path = defaultPath
	}
	if p.Cfg.Expiration == 0 {
		p.Cfg.Expiration = defaultExpiration
	}
	p.logger = log.New(os.Stderr, "prometheus_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		p.logger.SetOutput(logger.Writer())
		p.logger.SetFlags(logger.Flags())
	}
	p.replacer = strings.NewReplacer("-", "_", ":", "_", "/", "_")
	// create prometheus registery
	registry := prometheus.NewRegistry()

	err = registry.Register(p)
	if err != nil {
		return err
	}
	// create http server
	promHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError})

	mux := http.NewServeMux()
	mux.Handle(p.Cfg.Path, promHandler)

	p.server = &http.Server{
		Addr:    p.Cfg.Listen,
		Handler: mux,
	}

	// create tcp listener
	listener, err := net.Listen("tcp", p.Cfg.Listen)
	if err != nil {
		return err
	}
	// start worker
	p.wg.Add(2)
	go p.worker(ctx)
	go func() {
		defer p.wg.Done()
		err = p.server.Serve(listener)
		if err != nil && err != http.ErrServerClosed {
			p.logger.Printf("prometheus server error: %v", err)
		}
	}()
	p.logger.Printf("initialized prometheus output: %s", p.String())
	return nil
}
func (p *PrometheusOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			p.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case p.eventChan <- ev:
			}
		}
	}
}
func (p *PrometheusOutput) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := p.server.Shutdown(ctx)
	if err != nil {
		p.logger.Printf("failed to shutdown http server: %v", err)
	}
	p.cancelFn()
	p.wg.Wait()
	return nil
}
func (p *PrometheusOutput) Metrics() []prometheus.Collector { return p.metrics }

///
func (p *PrometheusOutput) Describe(ch chan<- *prometheus.Desc) {}

func (p *PrometheusOutput) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()
<commit_msg>rebase from master, use DecodeConfig method<commit_after>package prometheus_output

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	dto "github.com/prometheus/client_model/go"
	"google.golang.org/protobuf/proto"
)

const (
	defaultListen     = ":9804"
	defaultPath       = "/metrics"
	defaultExpiration = time.Minute
	defaultMetricHelp = "gNMIc generated metric"
)

type labelPair struct {
	Name  string
	Value string
}
type promMetric struct {
	name   string
	labels []*labelPair
	time   time.Time
	value  float64
}

func init() {
	outputs.Register("prometheus", func() outputs.Output {
		return &PrometheusOutput{
			Cfg:       &Config{},
			eventChan: make(chan *collector.EventMsg),
			wg:        new(sync.WaitGroup),
			entries:   make(map[uint64]*promMetric),
		}
	})
}

type PrometheusOutput struct {
	Cfg       *Config
	metrics   []prometheus.Collector
	logger    *log.Logger
	eventChan chan *collector.EventMsg

	wg     *sync.WaitGroup
	server *http.Server
	sync.Mutex
	entries map[uint64]*promMetric

	replacer *strings.Replacer
}
type Config struct {
	Listen     string        `mapstructure:"listen,omitempty"`
	Path       string        `mapstructure:"path,omitempty"`
	Expiration time.Duration `mapstructure:"expiration,omitempty"`
	Debug      bool          `mapstructure:"debug,omitempty"`
}

func (p *PrometheusOutput) String() string {
	b, err := json.Marshal(p)
	if err != nil {
		return ""
	}
	return string(b)
}
func (p *PrometheusOutput) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	err := outputs.DecodeConfig(cfg, p.Cfg)
	if err != nil {
		logger.Printf("prometheus output config decode failed: %v", err)
		return err
	}
	if p.Cfg.Listen == "" {
		p.Cfg.Listen = defaultListen
	}
	if p.Cfg.Path == "" {
		p.Cfg.Path = defaultPath
	}
	if p.Cfg.Expiration == 0 {
		p.Cfg.Expiration = defaultExpiration
	}
	p.logger = log.New(os.Stderr, "prometheus_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		p.logger.SetOutput(logger.Writer())
		p.logger.SetFlags(logger.Flags())
	}
	p.replacer = strings.NewReplacer("-", "_", ":", "_", "/", "_")
	// create prometheus registery
	registry := prometheus.NewRegistry()

	err = registry.Register(p)
	if err != nil {
		return err
	}
	// create http server
	promHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError})

	mux := http.NewServeMux()
	mux.Handle(p.Cfg.Path, promHandler)

	p.server = &http.Server{
		Addr:    p.Cfg.Listen,
		Handler: mux,
	}

	// create tcp listener
	listener, err := net.Listen("tcp", p.Cfg.Listen)
	if err != nil {
		return err
	}
	// start worker
	p.wg.Add(2)
	go p.worker(ctx)
	go func() {
		defer p.wg.Done()
		err = p.server.Serve(listener)
		if err != nil && err != http.ErrServerClosed {
			p.logger.Printf("prometheus server error: %v", err)
		}
	}()
	p.logger.Printf("initialized prometheus output: %s", p.String())
	go func() {
		<-ctx.Done()
		p.Close()
	}()
	return nil
}
func (p *PrometheusOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			p.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case p.eventChan <- ev:
			}
		}
	}
}
func (p *PrometheusOutput) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := p.server.Shutdown(ctx)
	if err != nil {
		p.logger.Printf("failed to shutdown http server: %v", err)
	}
	p.wg.Wait()
	return nil
}
func (p *PrometheusOutput) Metrics() []prometheus.Collector { return p.metrics }

///
func (p *PrometheusOutput) Describe(ch chan<- *prometheus.Desc) {}

func (p *PrometheusOutput) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()

{'fraction_non_alphanumeric': 0.14409993308052643, 'fraction_numerical': 0.0028998438545616775, 'mean_word_length': 3.7194736842105263, 'pattern_counts': {'":': 4, '<': 12, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3946138', 'n_tokens_mistral': 3449, 'n_tokens_neox': 3209, 'n_words': 707}