reuse OtelTraceState for ot value handling in the probability sampler

Signed-off-by: ayush-that <ayush1337@hotmail.com>

Author: ayush-that

sdk/src/trace/samplers/probability.cc

@@ -1,14 +1,11 @@
 // Copyright The OpenTelemetry Authors
 // SPDX-License-Identifier: Apache-2.0
 
-#include <stddef.h>
 #include <atomic>
 #include <cstdint>
 #include <map>
-#include <memory>
 #include <string>
 
-#include "opentelemetry/nostd/function_ref.h"
 #include "opentelemetry/nostd/shared_ptr.h"
 #include "opentelemetry/nostd/string_view.h"
 #include "opentelemetry/sdk/common/global_log_handler.h"
@@ -23,7 +20,6 @@
 #include "ot_trace_state.h"
 
 namespace trace_api = opentelemetry::trace;
-namespace nostd     = opentelemetry::nostd;
 
 namespace
 {
@@ -36,118 +32,13 @@ double ClampProbability(double ratio) noexcept
     return 1.0;
   return ratio;
 }
-
-bool ParseHex56(nostd::string_view hex, uint64_t *value) noexcept
-{
-  if (hex.size() != 14)
-    return false;
-  uint64_t result = 0;
-  for (char c : hex)
-  {
-    result <<= 4;
-    if (c >= '0' && c <= '9')
-      result |= static_cast<uint64_t>(c - '0');
-    else if (c >= 'a' && c <= 'f')
-      result |= static_cast<uint64_t>(c - 'a' + 10);
-    else
-      return false;
-  }
-  *value = result;
-  return true;
-}
-
-/**
- * Encodes a 56-bit threshold as lowercase hex with trailing zeros removed,
- * e.g. 2^55 becomes "8" and 0 becomes "0".
- */
-std::string EncodeThreshold(uint64_t threshold)
-{
-  if (threshold == 0)
-    return "0";
-  static const char kHex[] = "0123456789abcdef";
-  std::string hex(14, '0');
-  for (int i = 13; i >= 0; --i)
-  {
-    hex[i] = kHex[threshold & 0xf];
-    threshold >>= 4;
-  }
-  hex.erase(hex.find_last_not_of('0') + 1);
-  return hex;
-}
-
-// Returns true when the ot value carries a valid rv sub-key.
-bool ExplicitRandomness(nostd::string_view ot_value, uint64_t *randomness) noexcept
-{
-  while (!ot_value.empty())
-  {
-    size_t end                 = ot_value.find(';');
-    nostd::string_view sub_key = ot_value.substr(0, end);
-    ot_value                   = end == nostd::string_view::npos ? "" : ot_value.substr(end + 1);
-
-    if (sub_key.substr(0, 3) == "rv:")
-      return ParseHex56(sub_key.substr(3), randomness);
-  }
-  return false;
-}
-
-// Sub-keys that are empty, lack a key:value separator, or duplicate th are dropped.
-std::string SetThresholdSubKey(nostd::string_view ot_value, const std::string &threshold)
-{
-  std::string result;
-  bool replaced = false;
-  while (!ot_value.empty())
-  {
-    size_t end                 = ot_value.find(';');
-    nostd::string_view sub_key = ot_value.substr(0, end);
-    ot_value                   = end == nostd::string_view::npos ? "" : ot_value.substr(end + 1);
-
-    if (sub_key.empty() || sub_key.find(':') == nostd::string_view::npos)
-      continue;
-    if (sub_key.substr(0, 3) == "th:")
-    {
-      if (replaced)
-        continue;
-      replaced = true;
-      if (!result.empty())
-        result += ';';
-      result += "th:" + threshold;
-      continue;
-    }
-    if (!result.empty())
-      result += ';';
-    result.append(sub_key.data(), sub_key.size());
-  }
-  if (!replaced)
-  {
-    if (!result.empty())
-      result += ';';
-    result += "th:" + threshold;
-  }
-  return result;
-}
 }  // namespace
 
 OPENTELEMETRY_BEGIN_NAMESPACE
 namespace sdk
 {
 namespace trace
 {
-namespace
-{
-// Removes the inherited (now stale) "th" sub-key, keeping the other ot sub-keys.
-nostd::shared_ptr<trace_api::TraceState> EraseThreshold(
-    const nostd::shared_ptr<trace_api::TraceState> &parent_trace_state,
-    const std::string &ot_value)
-{
-  OtelTraceState ot_state = OtelTraceState::Parse(ot_value);
-  ot_state.has_threshold  = false;
-  std::string stripped    = ot_state.Serialize();
-  auto trace_state        = parent_trace_state->Delete(kOtTraceStateKey);
-  if (!stripped.empty())
-    trace_state = trace_state->Set(kOtTraceStateKey, stripped);
-  return trace_state;
-}
-}  // namespace
 
 ProbabilitySampler::ProbabilitySampler(double ratio)
     : description_("ProbabilitySampler{" + std::to_string(ClampProbability(ratio)) + "}"),
@@ -163,14 +54,26 @@ SamplingResult ProbabilitySampler::ShouldSample(
     const trace_api::SpanContextKeyValueIterable & /*links*/) noexcept
 {
   auto parent_trace_state = parent_context.trace_state();
+  if (!parent_trace_state)
+  {
+    parent_trace_state = trace_api::TraceState::GetDefault();
+  }
+
   std::string ot_value;
-  bool has_ot = parent_trace_state->Get(kOtTraceStateKey, ot_value);
+  parent_trace_state->Get(kOtTraceStateKey, ot_value);
+  OtelTraceState ot_state = OtelTraceState::Parse(ot_value);
 
   bool drop = threshold_ == kMaxThreshold;
   if (!drop)
   {
     uint64_t randomness = 0;
-    if (!ExplicitRandomness(ot_value, &randomness))
+    if (ot_state.has_random_value)
+    {
+      // An explicit "rv" from an upstream Level 2 participant wins over the
+      // trace id, keeping the sampling decision consistent across the trace.
+      randomness = ot_state.random_value;
+    }
+    else
     {
       if (parent_context.IsValid() && !parent_context.trace_flags().IsRandom())
       {
@@ -187,31 +90,27 @@ SamplingResult ProbabilitySampler::ShouldSample(
     drop = randomness < threshold_;
   }
 
+  // Record the effective threshold when sampling; a dropped span carries no
+  // probability, so its inherited (now stale) "th" must be erased. The "rv"
+  // sub-key and any other "ot" sub-keys are preserved by OtelTraceState.
   if (drop)
   {
-    return {Decision::DROP, nullptr, EraseThreshold(parent_trace_state, ot_value)};
+    ot_state.has_threshold = false;
   }
-
-  std::string new_ot_value = SetThresholdSubKey(ot_value, EncodeThreshold(threshold_));
-  if (!trace_api::TraceState::IsValidValue(new_ot_value))
-    return {Decision::RECORD_AND_SAMPLE, nullptr, EraseThreshold(parent_trace_state, ot_value)};
-
-  if (!has_ot)
+  else
   {
-    size_t entry_count = 0;
-    parent_trace_state->GetAllEntries([&entry_count](nostd::string_view, nostd::string_view) {
-      ++entry_count;
-      return true;
-    });
-    if (entry_count >= trace_api::TraceState::kMaxKeyValuePairs)
-      return {Decision::RECORD_AND_SAMPLE, nullptr, {}};
+    ot_state.has_threshold = true;
+    ot_state.threshold     = threshold_;
   }
 
-  auto trace_state =
-      has_ot ? parent_trace_state->Delete(kOtTraceStateKey)->Set(kOtTraceStateKey, new_ot_value)
-             : parent_trace_state->Set(kOtTraceStateKey, new_ot_value);
+  std::string new_ot_value = ot_state.Serialize();
+  auto trace_state         = parent_trace_state->Delete(kOtTraceStateKey);
+  if (!new_ot_value.empty())
+  {
+    trace_state = trace_state->Set(kOtTraceStateKey, new_ot_value);
+  }
 
-  return {Decision::RECORD_AND_SAMPLE, nullptr, trace_state};
+  return {drop ? Decision::DROP : Decision::RECORD_AND_SAMPLE, nullptr, trace_state};
 }
 
 nostd::string_view ProbabilitySampler::GetDescription() const noexcept

sdk/test/trace/probability_sampler_test.cc

@@ -177,7 +177,7 @@ TEST(ProbabilitySampler, ExplicitRandomnessTakesPrecedence)
 
   std::string ot_value;
   ASSERT_TRUE(sampling_result.trace_state->Get("ot", ot_value));
-  ASSERT_EQ("rv:ffffffffffffff;th:8", ot_value);
+  ASSERT_EQ("th:8;rv:ffffffffffffff", ot_value);
 
   // An rv of zero drops even when the trace id randomness would sample.
   trace_state = trace_api::TraceState::FromHeader("ot=rv:00000000000000");
@@ -302,7 +302,7 @@ TEST(ProbabilitySampler, MalformedSubKeysAreDropped)
   trace_api::SpanId span_id{span_id_buffer};
 
   // Empty, key-only, and duplicate th sub-keys are dropped from the output.
-  for (auto p : {std::make_pair("ot=a:1;;rv:ffffffffffffff", "a:1;rv:ffffffffffffff;th:8"),
+  for (auto p : {std::make_pair("ot=a:1;;rv:ffffffffffffff", "th:8;rv:ffffffffffffff;a:1"),
                  std::make_pair("ot=foo", "th:8"), std::make_pair("ot=th:1;th:2", "th:8")})
   {
     auto trace_state = trace_api::TraceState::FromHeader(p.first);
@@ -328,7 +328,7 @@ TEST(ProbabilitySampler, FullTraceStateKeepsParentTraceState)
   trace_api::SpanId span_id{span_id_buffer};
 
   // 32 entries and no ot entry: there is no room to add th, so the span is
-  // sampled but the trace state is left untouched.
+  // sampled and the parent entries are preserved without an ot entry.
   std::string header = "k0=0";
   for (int i = 1; i < 32; ++i)
   {
@@ -340,10 +340,17 @@ TEST(ProbabilitySampler, FullTraceStateKeepsParentTraceState)
   auto sampling_result = SampleWithContext(s1, context, TraceIdWithRandomness(0xffffffffffffff));
 
   ASSERT_EQ(Decision::RECORD_AND_SAMPLE, sampling_result.decision);
-  ASSERT_EQ(nullptr, sampling_result.trace_state);
+  ASSERT_NE(nullptr, sampling_result.trace_state);
+
+  std::string ot_value;
+  ASSERT_FALSE(sampling_result.trace_state->Get("ot", ot_value));
+
+  std::string k0_value;
+  ASSERT_TRUE(sampling_result.trace_state->Get("k0", k0_value));
+  ASSERT_EQ("0", k0_value);
 }
 
-TEST(ProbabilitySampler, OversizedValueKeepsParentTraceState)
+TEST(ProbabilitySampler, OversizedSubKeyDroppedToRecordThreshold)
 {
   ProbabilitySampler s1(0.5);
 
@@ -352,7 +359,8 @@ TEST(ProbabilitySampler, OversizedValueKeepsParentTraceState)
   uint8_t span_id_buffer[trace_api::SpanId::kSize] = {1};
   trace_api::SpanId span_id{span_id_buffer};
 
-  // The oversized ot value is re-emitted unchanged because it has no th to erase.
+  // A sampled span must record th; the foreign sub-key would push the ot value
+  // past the 256 char limit, so it is dropped to make room for th.
   std::string ot_value = "a:" + std::string(253, 'b');
   auto trace_state     = trace_api::TraceState::FromHeader("ot=" + ot_value);
   trace_api::SpanContext context(trace_id, span_id, trace_api::TraceFlags{0}, false, trace_state);
@@ -364,19 +372,20 @@ TEST(ProbabilitySampler, OversizedValueKeepsParentTraceState)
 
   std::string result_ot;
   ASSERT_TRUE(sampling_result.trace_state->Get("ot", result_ot));
-  ASSERT_EQ(ot_value, result_ot);
+  ASSERT_EQ("th:8", result_ot);
 }
 
-TEST(ProbabilitySampler, OversizedValueErasesStaleThreshold)
+TEST(ProbabilitySampler, OversizedSubKeyDroppedReplacesStaleThreshold)
 {
-  // ratio 0.1's 17-char th overflows the 256-char limit, so the stale th is erased.
   ProbabilitySampler s1(0.1);
 
   uint8_t trace_id_buffer[trace_api::TraceId::kSize] = {1};
   trace_api::TraceId trace_id{trace_id_buffer};
   uint8_t span_id_buffer[trace_api::SpanId::kSize] = {1};
   trace_api::SpanId span_id{span_id_buffer};
 
+  // The inherited th is replaced by this sampler's th; the foreign sub-key
+  // overflows the 256 char limit and is dropped, keeping the fresh th.
   std::string rest = "x:" + std::string(245, 'b');
   auto trace_state = trace_api::TraceState::FromHeader("ot=th:8;" + rest);
   trace_api::SpanContext context(trace_id, span_id, trace_api::TraceFlags{0}, false, trace_state);
@@ -388,8 +397,8 @@ TEST(ProbabilitySampler, OversizedValueErasesStaleThreshold)
 
   std::string result_ot;
   ASSERT_TRUE(sampling_result.trace_state->Get("ot", result_ot));
-  ASSERT_EQ(rest, result_ot);
-  ASSERT_EQ(std::string::npos, result_ot.find("th:"));
+  ASSERT_NE(std::string::npos, result_ot.find("th:"));
+  ASSERT_EQ(std::string::npos, result_ot.find("x:"));
 }
 
 TEST(ProbabilitySampler, IgnoresParentSampledFlag)