OTel-Go Consistent Probability Sampler and conformance tests (#1379)

* Second prototype (based on OTel-Go draft PR #2177 from main repo) * comments * add a test for sampled vs r/p-value consistency * adding tests * test various D values * test spec * have the test print the specification table * output the table used in the spec * remove a package * comment * add tracestate handling tests * remove a file * add parent test * 100% coverage * lint * propose short tests * clarify test settings * raise test timeout * raise more * disable statistical test w/ race detector * missing file * test fix * mod tidy * shorten the test runtime * even shorter * shorten the test; lint * run 3 / 15 non-deterministically * Update samplers/probability/consistent/sampler.go * Apply suggestions from code review Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com> * test one more thing * lint * changelog & readme * revert accidental golint fix * godoc Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com> Co-authored-by: Anthony Mirabella <a9@aneurysm9.com>
open-telemetry · Apr 26, 2022 · 5883a27 · 5883a27
1 parent 672ba73
commit 5883a27
Show file tree

Hide file tree

Showing 12 changed files with 1,680 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Added
+
+- Consistent probability sampler implementation. (#1379)
+
 ### Fixed
 
 - Fix the `otelmux` middleware by using `SpanKindServer` when deciding the `SpanStatus`.

diff --git a/samplers/probability/consistent/base2.go b/samplers/probability/consistent/base2.go
@@ -0,0 +1,69 @@
+// Copyright The OpenTelemetry Authors
+//
+// 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 consistent // import "go.opentelemetry.io/contrib/samplers/probability/consistent"
+
+import "math"
+
+// These are IEEE 754 double-width floating point constants used with
+// math.Float64bits.
+const (
+	offsetExponentMask = 0x7ff0000000000000
+	offsetExponentBias = 1023
+	significandBits    = 52
+)
+
+// expFromFloat64 returns floor(log2(x)).
+func expFromFloat64(x float64) int {
+	return int((math.Float64bits(x)&offsetExponentMask)>>significandBits) - offsetExponentBias
+}
+
+// expToFloat64 returns 2^x.
+func expToFloat64(x int) float64 {
+	return math.Float64frombits(uint64(offsetExponentBias+x) << significandBits)
+}
+
+// splitProb returns the two values of log-adjusted-count nearest to p
+// Example:
+//
+//   splitProb(0.375) => (2, 1, 0.5)
+//
+// indicates to sample with probability (2^-2) 50% of the time
+// and (2^-1) 50% of the time.
+func splitProb(p float64) (uint8, uint8, float64) {
+	if p < 2e-62 {
+		// Note: spec.
+		return pZeroValue, pZeroValue, 1
+	}
+	// Take the exponent and drop the significand to locate the
+	// smaller of two powers of two.
+	exp := expFromFloat64(p)
+
+	// Low is the smaller of two log-adjusted counts, the negative
+	// of the exponent computed above.
+	low := -exp
+	// High is the greater of two log-adjusted counts (i.e., one
+	// less than low, a smaller adjusted count means a larger
+	// probability).
+	high := low - 1
+
+	// Return these to probability values and use linear
+	// interpolation to compute the required probability of
+	// choosing the low-probability Sampler.
+	lowP := expToFloat64(-low)
+	highP := expToFloat64(-high)
+	lowProb := (highP - p) / (highP - lowP)
+
+	return uint8(low), uint8(high), lowProb
+}
diff --git a/samplers/probability/consistent/base2_test.go b/samplers/probability/consistent/base2_test.go
@@ -0,0 +1,61 @@
+// Copyright The OpenTelemetry Authors
+//
+// 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 consistent
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSplitProb(t *testing.T) {
+	require.Equal(t, -1, expFromFloat64(0.6))
+	require.Equal(t, -2, expFromFloat64(0.4))
+	require.Equal(t, 0.5, expToFloat64(-1))
+	require.Equal(t, 0.25, expToFloat64(-2))
+
+	for _, tc := range []struct {
+		in      float64
+		low     uint8
+		lowProb float64
+	}{
+		// Probability 0.75 corresponds with choosing S=1 (the
+		// "low" probability) 50% of the time and S=0 (the
+		// "high" probability) 50% of the time.
+		{0.75, 1, 0.5},
+		{0.6, 1, 0.8},
+		{0.9, 1, 0.2},
+
+		// Powers of 2 exactly
+		{1, 0, 1},
+		{0.5, 1, 1},
+		{0.25, 2, 1},
+
+		// Smaller numbers
+		{0.05, 5, 0.4},
+		{0.1, 4, 0.4}, // 0.1 == 0.4 * 1/16 + 0.6 * 1/8
+		{0.003, 9, 0.464},
+
+		// Special cases:
+		{0, 63, 1},
+	} {
+		low, high, lowProb := splitProb(tc.in)
+		require.Equal(t, tc.low, low, "got %v want %v", low, tc.low)
+		if lowProb != 1 {
+			require.Equal(t, tc.low-1, high, "got %v want %v", high, tc.low-1)
+		}
+		require.InEpsilon(t, tc.lowProb, lowProb, 1e-6, "got %v want %v", lowProb, tc.lowProb)
+	}
+}
diff --git a/samplers/probability/consistent/go.mod b/samplers/probability/consistent/go.mod
@@ -0,0 +1,10 @@
+module go.opentelemetry.io/contrib/samplers/probability/consistent
+
+go 1.16
+
+require (
+	github.com/stretchr/testify v1.7.1
+	go.opentelemetry.io/otel v1.6.1
+	go.opentelemetry.io/otel/sdk v1.6.1
+	go.opentelemetry.io/otel/trace v1.6.1
+)
diff --git a/samplers/probability/consistent/go.sum b/samplers/probability/consistent/go.sum
@@ -0,0 +1,27 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.2.3 h1:2DntVwHkVopvECVRSlL5PSo9eG+cAkDCuckLubN+rq0=
+github.com/go-logr/logr v1.2.3/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
+github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+go.opentelemetry.io/otel v1.6.1 h1:6r1YrcTenBvYa1x491d0GGpTVBsNECmrc/K6b+zDeis=
+go.opentelemetry.io/otel v1.6.1/go.mod h1:blzUabWHkX6LJewxvadmzafgh/wnvBSDBdOuwkAtrWQ=
+go.opentelemetry.io/otel/sdk v1.6.1 h1:ZmcNyMhcuAYIb/Nr6QhBPTMopMTbov/47wHt1gibkoY=
+go.opentelemetry.io/otel/sdk v1.6.1/go.mod h1:IVYrddmFZ+eJqu2k38qD3WezFR2pymCzm8tdxyh3R4E=
+go.opentelemetry.io/otel/trace v1.6.1 h1:f8c93l5tboBYZna1nWk0W9DYyMzJXDWdZcJZ0Kb400U=
+go.opentelemetry.io/otel/trace v1.6.1/go.mod h1:RkFRM1m0puWIq10oxImnGEduNBzxiN7TXluRBtE+5j0=
+golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7 h1:iGu644GcxtEcrInvDsQRCwJjtCIOlT2V7IRt6ah2Whw=
+golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/samplers/probability/consistent/parent.go b/samplers/probability/consistent/parent.go
@@ -0,0 +1,74 @@
+// Copyright The OpenTelemetry Authors
+//
+// 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 consistent // import "go.opentelemetry.io/contrib/samplers/probability/consistent"
+
+import (
+	"strings"
+
+	"go.opentelemetry.io/otel"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type (
+	parentProbabilitySampler struct {
+		delegate sdktrace.Sampler
+	}
+)
+
+// ParentProbabilityBased is an implementation of the OpenTelemetry
+// Trace Sampler interface that provides additional checks for tracestate
+// Probability Sampling fields.
+func ParentProbabilityBased(root sdktrace.Sampler, samplers ...sdktrace.ParentBasedSamplerOption) sdktrace.Sampler {
+	return &parentProbabilitySampler{
+		delegate: sdktrace.ParentBased(root, samplers...),
+	}
+}
+
+// ShouldSample implements "go.opentelemetry.io/otel/sdk/trace".Sampler.
+func (p *parentProbabilitySampler) ShouldSample(params sdktrace.SamplingParameters) sdktrace.SamplingResult {
+	psc := trace.SpanContextFromContext(params.ParentContext)
+
+	// Note: We do not check psc.IsValid(), i.e., we repair the tracestate
+	// with or without a parent TraceId and SpanId.
+	state := psc.TraceState()
+
+	otts, err := parseOTelTraceState(state.Get(traceStateKey), psc.IsSampled())
+
+	if err != nil {
+		otel.Handle(err)
+		value := otts.serialize()
+		if len(value) > 0 {
+			// Note: see the note in
+			// "go.opentelemetry.io/otel/trace".TraceState.Insert(). The
+			// error below is not a condition we're supposed to handle.
+			state, _ = state.Insert(traceStateKey, value)
+		} else {
+			state = state.Delete(traceStateKey)
+		}
+
+		// Fix the broken tracestate before calling the delegate.
+		params.ParentContext = trace.ContextWithSpanContext(params.ParentContext, psc.WithTraceState(state))
+	}
+
+	return p.delegate.ShouldSample(params)
+}
+
+// Description returns the same description as the built-in
+// ParentBased sampler, with "ParentBased" replaced by
+// "ParentProbabilityBased".
+func (p *parentProbabilitySampler) Description() string {
+	return "ParentProbabilityBased" + strings.TrimPrefix(p.delegate.Description(), "ParentBased")
+}