Add optional inner samplers to make ParentBasedSampler more customiza…

…ble (#1727) * WIP for ParentBasedSampler. Todo: finish comments, add tests for null samplers. * Better comments, tests for null arguments. * Remove unnecessary remark, left it there by mistake. * Make other samplers optional * Change to new delegating behavior. * Update CHANGELOG * Comment to note default values for inner samplers. Co-authored-by: Cijo Thomas <cithomas@microsoft.com>
open-telemetry · Feb 27, 2021 · 381908c · 381908c
1 parent 899a91b
commit 381908c
Show file tree

Hide file tree

Showing 7 changed files with 155 additions and 7 deletions.
diff --git a/src/OpenTelemetry/.publicApi/net452/PublicAPI.Unshipped.txt b/src/OpenTelemetry/.publicApi/net452/PublicAPI.Unshipped.txt
@@ -1,2 +1,3 @@
+OpenTelemetry.Trace.ParentBasedSampler.ParentBasedSampler(OpenTelemetry.Trace.Sampler rootSampler, OpenTelemetry.Trace.Sampler remoteParentSampled = null, OpenTelemetry.Trace.Sampler remoteParentNotSampled = null, OpenTelemetry.Trace.Sampler localParentSampled = null, OpenTelemetry.Trace.Sampler localParentNotSampled = null) -> void
 static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddLegacyActivity(this OpenTelemetry.Trace.TracerProviderBuilder tracerProviderBuilder, string operationName) -> OpenTelemetry.Trace.TracerProviderBuilder
 static OpenTelemetry.Trace.TracerProviderExtensions.ForceFlush(this OpenTelemetry.Trace.TracerProvider provider, int timeoutMilliseconds = -1) -> bool
diff --git a/src/OpenTelemetry/.publicApi/net46/PublicAPI.Unshipped.txt b/src/OpenTelemetry/.publicApi/net46/PublicAPI.Unshipped.txt
@@ -1,2 +1,3 @@
+OpenTelemetry.Trace.ParentBasedSampler.ParentBasedSampler(OpenTelemetry.Trace.Sampler rootSampler, OpenTelemetry.Trace.Sampler remoteParentSampled = null, OpenTelemetry.Trace.Sampler remoteParentNotSampled = null, OpenTelemetry.Trace.Sampler localParentSampled = null, OpenTelemetry.Trace.Sampler localParentNotSampled = null) -> void
 static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddLegacyActivity(this OpenTelemetry.Trace.TracerProviderBuilder tracerProviderBuilder, string operationName) -> OpenTelemetry.Trace.TracerProviderBuilder
 static OpenTelemetry.Trace.TracerProviderExtensions.ForceFlush(this OpenTelemetry.Trace.TracerProvider provider, int timeoutMilliseconds = -1) -> bool
diff --git a/src/OpenTelemetry/.publicApi/net461/PublicAPI.Unshipped.txt b/src/OpenTelemetry/.publicApi/net461/PublicAPI.Unshipped.txt
@@ -1,2 +1,3 @@
+OpenTelemetry.Trace.ParentBasedSampler.ParentBasedSampler(OpenTelemetry.Trace.Sampler rootSampler, OpenTelemetry.Trace.Sampler remoteParentSampled = null, OpenTelemetry.Trace.Sampler remoteParentNotSampled = null, OpenTelemetry.Trace.Sampler localParentSampled = null, OpenTelemetry.Trace.Sampler localParentNotSampled = null) -> void
 static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddLegacyActivity(this OpenTelemetry.Trace.TracerProviderBuilder tracerProviderBuilder, string operationName) -> OpenTelemetry.Trace.TracerProviderBuilder
 static OpenTelemetry.Trace.TracerProviderExtensions.ForceFlush(this OpenTelemetry.Trace.TracerProvider provider, int timeoutMilliseconds = -1) -> bool
diff --git a/src/OpenTelemetry/.publicApi/netstandard2.0/PublicAPI.Unshipped.txt b/src/OpenTelemetry/.publicApi/netstandard2.0/PublicAPI.Unshipped.txt
@@ -1,2 +1,3 @@
+OpenTelemetry.Trace.ParentBasedSampler.ParentBasedSampler(OpenTelemetry.Trace.Sampler rootSampler, OpenTelemetry.Trace.Sampler remoteParentSampled = null, OpenTelemetry.Trace.Sampler remoteParentNotSampled = null, OpenTelemetry.Trace.Sampler localParentSampled = null, OpenTelemetry.Trace.Sampler localParentNotSampled = null) -> void
 static OpenTelemetry.Trace.TracerProviderBuilderExtensions.AddLegacyActivity(this OpenTelemetry.Trace.TracerProviderBuilder tracerProviderBuilder, string operationName) -> OpenTelemetry.Trace.TracerProviderBuilder
 static OpenTelemetry.Trace.TracerProviderExtensions.ForceFlush(this OpenTelemetry.Trace.TracerProvider provider, int timeoutMilliseconds = -1) -> bool
diff --git a/src/OpenTelemetry/CHANGELOG.md b/src/OpenTelemetry/CHANGELOG.md
@@ -9,15 +9,17 @@ please check the latest changes
 
 ## Unreleased
 
-* Added `ForceFlush` to `TracerProvider`.
-  ([#1837](https://github.com/open-telemetry/opentelemetry-dotnet/pull/1837))
+* Added `ForceFlush` to `TracerProvider`. ([#1837](https://github.com/open-telemetry/opentelemetry-dotnet/pull/1837))
 
 * Added a TracerProvierBuilder extension method called
   `AddLegacyActivityOperationName` which is used by instrumentation libraries
   that use DiagnosticSource to get activities processed without
   ActivitySourceAdapter.
   [#1836](https://github.com/open-telemetry/opentelemetry-dotnet/pull/1836)
 
+* Added new constructor with optional parameters to allow customization of
+  `ParentBasedSampler` behavior. ([#1727](https://github.com/open-telemetry/opentelemetry-dotnet/pull/1727))
+
 ## 1.0.1
 
 Released 2021-Feb-10

diff --git a/src/OpenTelemetry/Trace/ParentBasedSampler.cs b/src/OpenTelemetry/Trace/ParentBasedSampler.cs
@@ -19,13 +19,23 @@
 namespace OpenTelemetry.Trace
 {
     /// <summary>
-    /// Sampler implementation which will take a sample if parent Activity or any linked Activity is sampled.
+    /// Sampler implementation which by default will take a sample if parent Activity or any linked Activity is sampled.
     /// Otherwise, samples root traces according to the specified root sampler.
     /// </summary>
+    /// <remarks>
+    /// The default behavior can be customized by providing additional samplers to be invoked for different
+    /// combinations of local/remote parent and its sampling decision.
+    /// See <see cref="ParentBasedSampler(Sampler, Sampler, Sampler, Sampler, Sampler)"/>.
+    /// </remarks>
     public sealed class ParentBasedSampler : Sampler
     {
         private readonly Sampler rootSampler;
 
+        private readonly Sampler remoteParentSampled;
+        private readonly Sampler remoteParentNotSampled;
+        private readonly Sampler localParentSampled;
+        private readonly Sampler localParentNotSampled;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ParentBasedSampler"/> class.
         /// </summary>
@@ -35,6 +45,50 @@ public ParentBasedSampler(Sampler rootSampler)
             this.rootSampler = rootSampler ?? throw new ArgumentNullException(nameof(rootSampler));
 
             this.Description = $"ParentBased{{{rootSampler.Description}}}";
+
+            this.remoteParentSampled = new AlwaysOnSampler();
+            this.remoteParentNotSampled = new AlwaysOffSampler();
+            this.localParentSampled = new AlwaysOnSampler();
+            this.localParentNotSampled = new AlwaysOffSampler();
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ParentBasedSampler"/> class with ability to delegate
+        /// sampling decision to one of the inner samplers provided.
+        /// </summary>
+        /// <param name="rootSampler">The <see cref="Sampler"/> to be called for root span/activity.</param>
+        /// <param name="remoteParentSampled">
+        /// A <see cref="Sampler"/> to delegate sampling decision to in case of
+        /// remote parent (<see cref="ActivityContext.IsRemote"/> == true) with <see cref="ActivityTraceFlags.Recorded"/> flag == true.
+        /// Default: <see cref="AlwaysOnSampler"/>.
+        /// </param>
+        /// <param name="remoteParentNotSampled">
+        /// A <see cref="Sampler"/> to delegate sampling decision to in case of
+        /// remote parent (<see cref="ActivityContext.IsRemote"/> == true) with <see cref="ActivityTraceFlags.Recorded"/> flag == false.
+        /// Default: <see cref="AlwaysOffSampler"/>.
+        /// </param>
+        /// <param name="localParentSampled">
+        /// A <see cref="Sampler"/> to delegate sampling decision to in case of
+        /// local parent (<see cref="ActivityContext.IsRemote"/> == false) with <see cref="ActivityTraceFlags.Recorded"/> flag == true.
+        /// Default: <see cref="AlwaysOnSampler"/>.
+        /// </param>
+        /// <param name="localParentNotSampled">
+        /// A <see cref="Sampler"/> to delegate sampling decision to in case of
+        /// local parent (<see cref="ActivityContext.IsRemote"/> == false) with <see cref="ActivityTraceFlags.Recorded"/> flag == false.
+        /// Default: <see cref="AlwaysOffSampler"/>.
+        /// </param>
+        public ParentBasedSampler(
+            Sampler rootSampler,
+            Sampler remoteParentSampled = null,
+            Sampler remoteParentNotSampled = null,
+            Sampler localParentSampled = null,
+            Sampler localParentNotSampled = null)
+            : this(rootSampler)
+        {
+            this.remoteParentSampled = remoteParentSampled ?? new AlwaysOnSampler();
+            this.remoteParentNotSampled = remoteParentNotSampled ?? new AlwaysOffSampler();
+            this.localParentSampled = localParentSampled ?? new AlwaysOnSampler();
+            this.localParentNotSampled = localParentNotSampled ?? new AlwaysOffSampler();
         }
 
         /// <inheritdoc />
@@ -47,10 +101,17 @@ public override SamplingResult ShouldSample(in SamplingParameters samplingParame
                 return this.rootSampler.ShouldSample(samplingParameters);
             }
 
-            // If the parent is sampled keep the sampling decision.
+            // Is parent sampled?
             if ((parentContext.TraceFlags & ActivityTraceFlags.Recorded) != 0)
             {
-                return new SamplingResult(SamplingDecision.RecordAndSample);
+                if (parentContext.IsRemote)
+                {
+                    return this.remoteParentSampled.ShouldSample(samplingParameters);
+                }
+                else
+                {
+                    return this.localParentSampled.ShouldSample(samplingParameters);
+                }
             }
 
             if (samplingParameters.Links != null)
@@ -68,8 +129,16 @@ public override SamplingResult ShouldSample(in SamplingParameters samplingParame
                 }
             }
 
-            // If parent was not sampled, do not sample.
-            return new SamplingResult(SamplingDecision.Drop);
+            // If parent was not sampled (and no linked context exists) => delegate to the "not sampled"
+            // inner samplers.
+            if (parentContext.IsRemote)
+            {
+                return this.remoteParentNotSampled.ShouldSample(samplingParameters);
+            }
+            else
+            {
+                return this.localParentNotSampled.ShouldSample(samplingParameters);
+            }
         }
     }
 }
diff --git a/test/OpenTelemetry.Tests/Trace/ParentBasedSamplerTests.cs b/test/OpenTelemetry.Tests/Trace/ParentBasedSamplerTests.cs
@@ -13,7 +13,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // </copyright>
+
+using System;
 using System.Diagnostics;
+using Moq;
 using Xunit;
 
 namespace OpenTelemetry.Trace.Tests
@@ -111,5 +114,75 @@ public void SampledParentLink()
                         kind: ActivityKind.Client,
                         links: sampledLink)));
         }
+
+        [Theory]
+        [InlineData(true, true)]
+        [InlineData(true, false)]
+        [InlineData(false, true)]
+        [InlineData(false, false)]
+        public void CustomSamplers(bool parentIsRemote, bool parentIsSampled)
+        {
+            var mockRepository = new MockRepository(MockBehavior.Strict);
+            var remoteParentSampled = mockRepository.Create<Sampler>();
+            var remoteParentNotSampled = mockRepository.Create<Sampler>();
+            var localParentSampled = mockRepository.Create<Sampler>();
+            var localParentNotSampled = mockRepository.Create<Sampler>();
+
+            var samplerUnderTest = new ParentBasedSampler(
+                new AlwaysOnSampler(), // root
+                remoteParentSampled.Object,
+                remoteParentNotSampled.Object,
+                localParentSampled.Object,
+                localParentNotSampled.Object);
+
+            var samplingParams = this.MakeTestParameters(parentIsRemote, parentIsSampled);
+
+            Mock<Sampler> invokedSampler;
+            if (parentIsRemote && parentIsSampled)
+            {
+                invokedSampler = remoteParentSampled;
+            }
+            else if (parentIsRemote && !parentIsSampled)
+            {
+                invokedSampler = remoteParentNotSampled;
+            }
+            else if (!parentIsRemote && parentIsSampled)
+            {
+                invokedSampler = localParentSampled;
+            }
+            else
+            {
+                invokedSampler = localParentNotSampled;
+            }
+
+            var expectedResult = new SamplingResult(SamplingDecision.RecordAndSample);
+            invokedSampler.Setup(sampler => sampler.ShouldSample(samplingParams)).Returns(expectedResult);
+
+            var actualResult = samplerUnderTest.ShouldSample(samplingParams);
+
+            mockRepository.VerifyAll();
+            Assert.Equal(expectedResult, actualResult);
+            mockRepository.VerifyNoOtherCalls();
+        }
+
+        [Fact]
+        public void DisallowNullRootSampler()
+        {
+            Assert.Throws<ArgumentNullException>(() => new ParentBasedSampler(null));
+        }
+
+        private SamplingParameters MakeTestParameters(bool parentIsRemote, bool parentIsSampled)
+        {
+            return new SamplingParameters(
+                parentContext: new ActivityContext(
+                    ActivityTraceId.CreateRandom(),
+                    ActivitySpanId.CreateRandom(),
+                    parentIsSampled ? ActivityTraceFlags.Recorded : ActivityTraceFlags.None,
+                    null,
+                    parentIsRemote),
+                traceId: default,
+                name: "Span",
+                kind: ActivityKind.Client);
+        }
     }
 }