<commit_before>package samplers

import (
	"bytes"
	"errors"
	"fmt"
	"hash/fnv"
	"math"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/stripe/veneur/ssf"
)

var invalidMetricTypeError = errors.New("Invalid type for metric")

// UDPMetric is a representation of the sample provided by a client. The tag list
// should be deterministically ordered.
type UDPMetric struct {
	MetricKey
	Digest     uint32
	Value      interface{}
	SampleRate float32
	Tags       []string
	Scope      MetricScope
}

type MetricScope int

const (
	MixedScope MetricScope = iota
	LocalOnly
	GlobalOnly
)

// MetricKey is a struct used to key the metrics into the worker's map. All fields must be comparable types.
type MetricKey struct {
	Name       string `json:"name"`
	Type       string `json:"type"`
	JoinedTags string `json:"tagstring"` // tags in deterministic order, joined with commas
}

// ToString returns a string representation of this MetricKey
func (m *MetricKey) String() string {
	var buff bytes.Buffer
	buff.WriteString(m.Name)
	buff.WriteString(m.Type)
	buff.WriteString(m.JoinedTags)
	return buff.String()
}

// ValidTrace takes in an SSF span and determines if it is valid or not.
// It also makes sure the Tags is non-nil, since we use it later.
func ValidTrace(sample *ssf.SSFSpan) bool {
	ret := true
	ret = ret && sample.Id != 0
	ret = ret && sample.TraceId != 0
	ret = ret && sample.StartTimestamp != 0
	ret = ret && sample.EndTimestamp != 0
	return ret
}

// ValidMetric takes in an SSF sample and determines if it is valid or not.
func ValidMetric(sample UDPMetric) bool {
	ret := true
	ret = ret && sample.Name != ""
	ret = ret && sample.Value != nil
	return ret
}

// A Message struct represents a parsed SSF message. It encapsulates
// an ssf.SSFSpan object, which can contain a trace span and/or a set
// of metrics.
type Message struct {
	span *ssf.SSFSpan
}

// InvalidTrace is an error type indicating that an SSF span was
// invalid.
type InvalidTrace struct {
	span *ssf.SSFSpan
}

func (e *InvalidTrace) Error() string {
	return fmt.Sprintf("not a valid trace span: %#v", e.span)
}

// TraceSpan checks if an SSF message is a valid trace. If so, it
// returns a pointer to that original span. If the span is not a valid
// trace, TraceSpan returns nil and an *InvalidTrace error type.
//
// The span returned from TraceSpan does contain the (unparsed) metric
// samples contained in the span. Note also that since the data
// returned is a pointer to the original span, it's not advisable to
// modify that span directly.
func (m *Message) TraceSpan() (*ssf.SSFSpan, error) {
	if !ValidTrace(m.span) {
		return nil, &InvalidTrace{m.span}
	}
	return m.span, nil
}

// Metrics examines an SSF message, parses and returns a new array
// containing any metrics contained in the message. If any parse error
// occurs in processing any of the metrics, ExtractMetrics collects
// them into the error type InvalidMetrics and returns this error
// alongside any valid metrics that could be parsed.
func (m *Message) Metrics() ([]UDPMetric, error) {
	span := m.span
	invalid := []*ssf.SSFSample{}
	metrics := make([]UDPMetric, 0, len(span.Metrics))
	for _, metricPacket := range span.Metrics {
		metric, err := ParseMetricSSF(metricPacket)
		if err != nil || !ValidMetric(metric) {
			invalid = append(invalid, metricPacket)
			continue
		}
		metrics = append(metrics, metric)
	}
	if len(invalid) != 0 {
		return metrics, &invalidMetrics{invalid}
	}
	return metrics, nil
}

var pbufPool = sync.Pool{
	New: func() interface{} {
		return proto.NewBuffer(nil)
	},
}

// ParseSSF takes in a byte slice and returns: a normalized SSFSpan
// and an error if any errors in parsing the SSF packet occur.
func ParseSSF(packet []byte) (*Message, error) {
	span := &ssf.SSFSpan{}
	scratchBuff := pbufPool.Get().(*proto.Buffer)
	defer func() {
		scratchBuff.Reset()
		pbufPool.Put(scratchBuff)
	}()
	scratchBuff.SetBuf(packet)
	err := scratchBuff.Unmarshal(span)

	if err != nil {
		return nil, err
	}

	// Normalize the span:
	if span.Tags == nil {
		span.Tags = make(map[string]string, 0)
	}
	if span.Name == "" {
		// Even though incoming packets should have Name set,
		// this allows Veneur to be backwards-compatible.
		for k, v := range span.Tags {
			if k == "name" {
				span.Name = v
			}
		}
		delete(span.Tags, "name")
	}

	// Normalize metrics on the span:
	for _, sample := range span.Metrics {
		if sample.SampleRate == 0 {
			sample.SampleRate = 1
		}
	}
	return &Message{span}, nil
}

// InvalidMetrics is an error type returned if any metric could not be parsed.
type InvalidMetrics interface {
	error

	// Samples returns any samples that couldn't be parsed or validated.
	Samples() []*ssf.SSFSample
}
<commit_msg>Initialize the Tags field better<commit_after>package samplers

import (
	"bytes"
	"errors"
	"fmt"
	"hash/fnv"
	"math"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/stripe/veneur/ssf"
)

var invalidMetricTypeError = errors.New("Invalid type for metric")

// UDPMetric is a representation of the sample provided by a client. The tag list
// should be deterministically ordered.
type UDPMetric struct {
	MetricKey
	Digest     uint32
	Value      interface{}
	SampleRate float32
	Tags       []string
	Scope      MetricScope
}

type MetricScope int

const (
	MixedScope MetricScope = iota
	LocalOnly
	GlobalOnly
)

// MetricKey is a struct used to key the metrics into the worker's map. All fields must be comparable types.
type MetricKey struct {
	Name       string `json:"name"`
	Type       string `json:"type"`
	JoinedTags string `json:"tagstring"` // tags in deterministic order, joined with commas
}

// ToString returns a string representation of this MetricKey
func (m *MetricKey) String() string {
	var buff bytes.Buffer
	buff.WriteString(m.Name)
	buff.WriteString(m.Type)
	buff.WriteString(m.JoinedTags)
	return buff.String()
}

// ValidTrace takes in an SSF span and determines if it is valid or not.
// It also makes sure the Tags is non-nil, since we use it later.
func ValidTrace(sample *ssf.SSFSpan) bool {
	ret := true
	ret = ret && sample.Id != 0
	ret = ret && sample.TraceId != 0
	ret = ret && sample.StartTimestamp != 0
	ret = ret && sample.EndTimestamp != 0
	return ret
}

// ValidMetric takes in an SSF sample and determines if it is valid or not.
func ValidMetric(sample UDPMetric) bool {
	ret := true
	ret = ret && sample.Name != ""
	ret = ret && sample.Value != nil
	return ret
}

// A Message struct represents a parsed SSF message. It encapsulates
// an ssf.SSFSpan object, which can contain a trace span and/or a set
// of metrics.
type Message struct {
	span *ssf.SSFSpan
}

// InvalidTrace is an error type indicating that an SSF span was
// invalid.
type InvalidTrace struct {
	span *ssf.SSFSpan
}

func (e *InvalidTrace) Error() string {
	return fmt.Sprintf("not a valid trace span: %#v", e.span)
}

// TraceSpan checks if an SSF message is a valid trace. If so, it
// returns a pointer to that original span. If the span is not a valid
// trace, TraceSpan returns nil and an *InvalidTrace error type.
//
// The span returned from TraceSpan does contain the (unparsed) metric
// samples contained in the span. Note also that since the data
// returned is a pointer to the original span, it's not advisable to
// modify that span directly.
func (m *Message) TraceSpan() (*ssf.SSFSpan, error) {
	if !ValidTrace(m.span) {
		return nil, &InvalidTrace{m.span}
	}
	return m.span, nil
}

// Metrics examines an SSF message, parses and returns a new array
// containing any metrics contained in the message. If any parse error
// occurs in processing any of the metrics, ExtractMetrics collects
// them into the error type InvalidMetrics and returns this error
// alongside any valid metrics that could be parsed.
func (m *Message) Metrics() ([]UDPMetric, error) {
	span := m.span
	invalid := []*ssf.SSFSample{}
	metrics := make([]UDPMetric, 0, len(span.Metrics))
	for _, metricPacket := range span.Metrics {
		metric, err := ParseMetricSSF(metricPacket)
		if err != nil || !ValidMetric(metric) {
			invalid = append(invalid, metricPacket)
			continue
		}
		metrics = append(metrics, metric)
	}
	if len(invalid) != 0 {
		return metrics, &invalidMetrics{invalid}
	}
	return metrics, nil
}

var pbufPool = sync.Pool{
	New: func() interface{} {
		return proto.NewBuffer(nil)
	},
}

// ParseSSF takes in a byte slice and returns: a normalized SSFSpan
// and an error if any errors in parsing the SSF packet occur.
func ParseSSF(packet []byte) (*Message, error) {
	span := &ssf.SSFSpan{}
	scratchBuff := pbufPool.Get().(*proto.Buffer)
	defer func() {
		scratchBuff.Reset()
		pbufPool.Put(scratchBuff)
	}()
	scratchBuff.SetBuf(packet)
	err := scratchBuff.Unmarshal(span)

	if err != nil {
		return nil, err
	}

	// Normalize the span:
	if span.Tags == nil {
		span.Tags = map[string]string{}
	}
	if span.Name == "" {
		// Even though incoming packets should have Name set,
		// this allows Veneur to be backwards-compatible.
		for k, v := range span.Tags {
			if k == "name" {
				span.Name = v
			}
		}
		delete(span.Tags, "name")
	}

	// Normalize metrics on the span:
	for _, sample := range span.Metrics {
		if sample.SampleRate == 0 {
			sample.SampleRate = 1
		}
	}
	return &Message{span}, nil
}

// InvalidMetrics is an error type returned if any metric could not be parsed.
type InvalidMetrics interface {
	error

	// Samples returns any samples that couldn't be parsed or validated.
	Samples() []*ssf.SSFSample
}

{'fraction_non_alphanumeric': 0.10649214659685864, 'fraction_numerical': 0.002617801047120419, 'mean_word_length': 3.5567748091603053, 'pattern_counts': {'":': 0, '<': 3, '<?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': '1412469', 'n_tokens_mistral': 3127, 'n_tokens_neox': 3014, 'n_words': 1149}