field_sampler.go

package fieldsampler

import (
	"github.com/honeycombio/beeline-go/sample"
)

const (
	// CustomSampleRateFieldName is the field to use to override the default Sample Rate
	CustomSampleRateFieldName = "CustomSampleRate"

	// DeterministicSampleFieldName is the field used as a determinant for the sampler
	DeterministicSampleFieldName = "trace.trace_id"
)

// NewFieldSamplerHook creates a new SamplerHook that respects CustomSampleRateFieldName if present,
// otherwise it uses the given defaultSampleRate. Then it delegates to the beeline deterministic sampler.
func NewFieldSamplerHook(defaultSampleRate uint) (func(map[string]interface{}) (bool, int), error) {
	defaultSampler, err := sample.NewDeterministicSampler(defaultSampleRate)
	if err != nil {
		return nil, err
	}

	return func(fields map[string]interface{}) (bool, int) {
		sampler := defaultSampler

		rate := int(defaultSampleRate)
		if custom, ok := fields[CustomSampleRateFieldName].(int); ok {
			rate = custom
			newSampler, err := sample.NewDeterministicSampler(uint(rate))
			if err == nil && rate > 0 {
				sampler = newSampler
			}
		}

		determinant := ""
		if traceID, ok := fields[DeterministicSampleFieldName].(string); ok {
			determinant = traceID
		}

		shouldSample := sampler.Sample(determinant)
		return shouldSample, rate
	}, nil
}