From a40750d502423af3309835078428be18f72057c4 Mon Sep 17 00:00:00 2001
From: gillesbergerp <paulgillesberger@live.com>
Date: Sat, 11 Nov 2023 05:32:09 +0100
Subject: [PATCH] Implement a hierarchical sampler

---
 .../opentelemetry/sampling/hierarchical.rb    |   1 +
 .../sampling/hierarchical/sampler.rb          |  37 ++++
 .../sampling/hierarchical/sampler_test.rb     | 168 ++++++++++++++++++
 .../test/opentelemetry/test_helper.rb         |  17 ++
 sampling/hierarchical/test/test_helper.rb     |   1 -
 5 files changed, 223 insertions(+), 1 deletion(-)
 create mode 100644 sampling/hierarchical/lib/opentelemetry/sampling/hierarchical/sampler.rb
 create mode 100644 sampling/hierarchical/test/opentelemetry/sampling/hierarchical/sampler_test.rb
 create mode 100644 sampling/hierarchical/test/opentelemetry/test_helper.rb

diff --git a/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical.rb b/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical.rb
index 40fce77778..cd71fe0f8a 100644
--- a/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical.rb
+++ b/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical.rb
@@ -11,4 +11,5 @@
 # The OpenTelemetry module provides global accessors for telemetry objects.
 # See the documentation for the `opentelemetry-api` gem for details.
 
+require_relative('hierarchical/sampler')
 require_relative('hierarchical/version')
diff --git a/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical/sampler.rb b/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical/sampler.rb
new file mode 100644
index 0000000000..f5e8ddd9d4
--- /dev/null
+++ b/sampling/hierarchical/lib/opentelemetry/sampling/hierarchical/sampler.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module OpenTelemetry
+  module Sampling
+    module Hierarchical
+      class Sampler
+        # @param [Array<OpenTelemetry::SDK::Trace::Samplers>] samplers
+        def initialize(*samplers)
+          @samplers = samplers
+        end
+
+        # @param [String] trace_id
+        # @param [OpenTelemetry::Context] parent_context
+        # @param [Enumerable<Link>] links
+        # @param [String] name
+        # @param [Symbol] kind
+        # @param [Hash<String, Object>] attributes
+        # @return [OpenTelemetry::SDK::Trace::Samplers::Result] The sampling result.
+        def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
+          @samplers.each do |sampler|
+            result = sampler.should_sample?(trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes)
+            return result if result.sampled?
+          end
+
+          OpenTelemetry::SDK::Trace::Samplers::Result.new(
+            decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
+            tracestate: OpenTelemetry::Trace.current_span(parent_context).context.tracestate
+          )
+        end
+      end
+    end
+  end
+end
diff --git a/sampling/hierarchical/test/opentelemetry/sampling/hierarchical/sampler_test.rb b/sampling/hierarchical/test/opentelemetry/sampling/hierarchical/sampler_test.rb
new file mode 100644
index 0000000000..a7df163524
--- /dev/null
+++ b/sampling/hierarchical/test/opentelemetry/sampling/hierarchical/sampler_test.rb
@@ -0,0 +1,168 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require('test_helper')
+
+describe(OpenTelemetry::Sampling::Hierarchical::Sampler) do
+  describe('#should_sample?') do
+    it('returns DROP if it has no samplers') do
+      sampler = OpenTelemetry::Sampling::Hierarchical::Sampler.new([])
+
+      _(
+        sampler.should_sample?(
+          trace_id: SecureRandom.uuid.to_s,
+          parent_context: nil,
+          links: [],
+          name: SecureRandom.uuid.to_s,
+          kind: :internal,
+          attributes: {}
+        ).sampled?
+      ).must_equal(false)
+    end
+
+    it('returns RECORD_AND_SAMPLE if the first sampler returns RECORD_AND_SAMPLE') do
+      trace_id = SecureRandom.uuid.to_s
+      name = SecureRandom.uuid.to_s
+
+      first_sampler = Minitest::Mock.new
+      second_sampler = Minitest::Mock.new
+
+      first_sampler.expect(
+        :should_sample?,
+        OpenTelemetry::SDK::Trace::Samplers::Result.new(
+          decision: OpenTelemetry::SDK::Trace::Samplers::Decision::RECORD_AND_SAMPLE,
+          tracestate: OpenTelemetry::Trace::Tracestate.from_hash({})
+        ),
+        [],
+        trace_id: trace_id,
+        parent_context: nil,
+        links: [],
+        name: name,
+        kind: :internal,
+        attributes: {}
+      )
+
+      sampler = OpenTelemetry::Sampling::Hierarchical::Sampler.new([first_sampler, second_sampler])
+      _(
+        sampler.should_sample?(
+          trace_id: trace_id,
+          parent_context: nil,
+          links: [],
+          name: name,
+          kind: :internal,
+          attributes: {}
+        ).sampled?
+      ).must_equal(true)
+
+      first_sampler.verify
+      second_sampler.verify
+    end
+
+    it('returns RECORD_AND_SAMPLE if the second sampler returns RECORD_AND_SAMPLE') do
+      trace_id = SecureRandom.uuid.to_s
+      name = SecureRandom.uuid.to_s
+
+      first_sampler = Minitest::Mock.new
+      second_sampler = Minitest::Mock.new
+
+      first_sampler.expect(
+        :should_sample?,
+        OpenTelemetry::SDK::Trace::Samplers::Result.new(
+          decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
+          tracestate: OpenTelemetry::Trace::Tracestate.from_hash({})
+        ),
+        [],
+        trace_id: trace_id,
+        parent_context: nil,
+        links: [],
+        name: name,
+        kind: :internal,
+        attributes: {}
+      )
+      second_sampler.expect(
+        :should_sample?,
+        OpenTelemetry::SDK::Trace::Samplers::Result.new(
+          decision: OpenTelemetry::SDK::Trace::Samplers::Decision::RECORD_AND_SAMPLE,
+          tracestate: OpenTelemetry::Trace::Tracestate.from_hash({})
+        ),
+        [],
+        trace_id: trace_id,
+        parent_context: nil,
+        links: [],
+        name: name,
+        kind: :internal,
+        attributes: {}
+      )
+
+      sampler = OpenTelemetry::Sampling::Hierarchical::Sampler.new([first_sampler, second_sampler])
+      _(
+        sampler.should_sample?(
+          trace_id: trace_id,
+          parent_context: nil,
+          links: [],
+          name: name,
+          kind: :internal,
+          attributes: {}
+        ).sampled?
+      ).must_equal(true)
+
+      first_sampler.verify
+      second_sampler.verify
+    end
+
+    it('returns DROP if both samplers return DROP') do
+      trace_id = SecureRandom.uuid.to_s
+      name = SecureRandom.uuid.to_s
+
+      first_sampler = Minitest::Mock.new
+      second_sampler = Minitest::Mock.new
+
+      first_sampler.expect(
+        :should_sample?,
+        OpenTelemetry::SDK::Trace::Samplers::Result.new(
+          decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
+          tracestate: OpenTelemetry::Trace::Tracestate.from_hash({})
+        ),
+        [],
+        trace_id: trace_id,
+        parent_context: nil,
+        links: [],
+        name: name,
+        kind: :internal,
+        attributes: {}
+      )
+      second_sampler.expect(
+        :should_sample?,
+        OpenTelemetry::SDK::Trace::Samplers::Result.new(
+          decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
+          tracestate: OpenTelemetry::Trace::Tracestate.from_hash({})
+        ),
+        [],
+        trace_id: trace_id,
+        parent_context: nil,
+        links: [],
+        name: name,
+        kind: :internal,
+        attributes: {}
+      )
+
+      sampler = OpenTelemetry::Sampling::Hierarchical::Sampler.new([first_sampler, second_sampler])
+      _(
+        sampler.should_sample?(
+          trace_id: trace_id,
+          parent_context: nil,
+          links: [],
+          name: name,
+          kind: :internal,
+          attributes: {}
+        ).sampled?
+      ).must_equal(false)
+
+      first_sampler.verify
+      second_sampler.verify
+    end
+  end
+end
diff --git a/sampling/hierarchical/test/opentelemetry/test_helper.rb b/sampling/hierarchical/test/opentelemetry/test_helper.rb
new file mode 100644
index 0000000000..45157c0388
--- /dev/null
+++ b/sampling/hierarchical/test/opentelemetry/test_helper.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require('bundler/setup')
+Bundler.require(:default, :development, :test)
+
+SimpleCov.minimum_coverage(85)
+SimpleCov.start
+
+require('opentelemetry-sampling-hierarchical')
+require('minitest/autorun')
+require('webmock/minitest')
+
+OpenTelemetry.logger = Logger.new($stderr, level: ENV.fetch('OTEL_LOG_LEVEL', 'error').to_sym)
diff --git a/sampling/hierarchical/test/test_helper.rb b/sampling/hierarchical/test/test_helper.rb
index 54d49cc81a..ed80951746 100644
--- a/sampling/hierarchical/test/test_helper.rb
+++ b/sampling/hierarchical/test/test_helper.rb
@@ -12,6 +12,5 @@
 
 require('opentelemetry-sampling-hierarchical')
 require('minitest/autorun')
-require('webmock/minitest')
 
 OpenTelemetry.logger = Logger.new($stderr, level: ENV.fetch('OTEL_LOG_LEVEL', 'debug').to_sym)