1// Licensed to Elasticsearch B.V. under one or more contributor
2// license agreements. See the NOTICE file distributed with
3// this work for additional information regarding copyright
4// ownership. Elasticsearch B.V. licenses this file to you under
5// the Apache License, Version 2.0 (the "License"); you may
6// not use this file except in compliance with the License.
7// You may obtain a copy of the License at
8//
9//     http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18package apm_test
19
20import (
21	"encoding/binary"
22	"math/rand"
23	"sync"
24	"testing"
25
26	"github.com/stretchr/testify/assert"
27
28	"go.elastic.co/apm"
29)
30
31func TestRatioSampler(t *testing.T) {
32	ratio := 0.75
33	s := apm.NewRatioSampler(ratio)
34
35	const (
36		numGoroutines = 100
37		numIterations = 1000
38	)
39
40	sampled := make([]int, numGoroutines)
41	var wg sync.WaitGroup
42	for i := 0; i < numGoroutines; i++ {
43		wg.Add(1)
44		go func(i int) {
45			defer wg.Done()
46
47			// fixed seed to avoid intermittent failures
48			rng := rand.New(rand.NewSource(int64(i)))
49			for j := 0; j < numIterations; j++ {
50				var traceContext apm.TraceContext
51				binary.LittleEndian.PutUint64(traceContext.Span[:], rng.Uint64())
52				if s.Sample(traceContext) {
53					sampled[i]++
54				}
55			}
56
57		}(i)
58	}
59	wg.Wait()
60
61	var total int
62	for i := 0; i < numGoroutines; i++ {
63		total += sampled[i]
64	}
65	assert.InDelta(t, ratio, float64(total)/(numGoroutines*numIterations), 0.1)
66}
67
68func TestRatioSamplerAlways(t *testing.T) {
69	s := apm.NewRatioSampler(1.0)
70	assert.False(t, s.Sample(apm.TraceContext{})) // invalid span ID
71	assert.True(t, s.Sample(apm.TraceContext{
72		Span: apm.SpanID{0, 0, 0, 0, 0, 0, 0, 1},
73	}))
74	assert.True(t, s.Sample(apm.TraceContext{
75		Span: apm.SpanID{255, 255, 255, 255, 255, 255, 255, 255},
76	}))
77}
78
79func TestRatioSamplerNever(t *testing.T) {
80	s := apm.NewRatioSampler(0)
81	assert.False(t, s.Sample(apm.TraceContext{})) // invalid span ID
82	assert.False(t, s.Sample(apm.TraceContext{
83		Span: apm.SpanID{0, 0, 0, 0, 0, 0, 0, 1},
84	}))
85	assert.False(t, s.Sample(apm.TraceContext{
86		Span: apm.SpanID{255, 255, 255, 255, 255, 255, 255, 255},
87	}))
88}
89