strip stale th on the invalid-ot sample path

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

Author: ayush-that

sdk/src/trace/samplers/probability.cc

@@ -9,6 +9,7 @@
 #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"
 #include "opentelemetry/sdk/trace/sampler.h"
@@ -131,6 +132,23 @@ 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)) + "}"),
       threshold_(CalculateThreshold(ClampProbability(ratio)))
@@ -171,18 +189,12 @@ SamplingResult ProbabilitySampler::ShouldSample(
 
   if (drop)
   {
-    OtelTraceState ot_state = OtelTraceState::Parse(ot_value);
-    ot_state.has_threshold  = false;
-    std::string dropped_ot  = ot_state.Serialize();
-    auto trace_state        = parent_trace_state->Delete(kOtTraceStateKey);
-    if (!dropped_ot.empty())
-      trace_state = trace_state->Set(kOtTraceStateKey, dropped_ot);
-    return {Decision::DROP, nullptr, trace_state};
+    return {Decision::DROP, nullptr, EraseThreshold(parent_trace_state, ot_value)};
   }
 
   std::string new_ot_value = SetThresholdSubKey(ot_value, EncodeThreshold(threshold_));
   if (!trace_api::TraceState::IsValidValue(new_ot_value))
-    return {Decision::RECORD_AND_SAMPLE, nullptr, {}};
+    return {Decision::RECORD_AND_SAMPLE, nullptr, EraseThreshold(parent_trace_state, ot_value)};
 
   if (!has_ot)
   {

sdk/test/trace/probability_sampler_test.cc

@@ -352,16 +352,44 @@ TEST(ProbabilitySampler, OversizedValueKeepsParentTraceState)
   uint8_t span_id_buffer[trace_api::SpanId::kSize] = {1};
   trace_api::SpanId span_id{span_id_buffer};
 
-  // Appending the th sub-key would exceed the 256 character value limit, so
-  // the span is sampled but the trace state is left untouched.
+  // The oversized ot value is re-emitted unchanged because it has no th to erase.
   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);
 
   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 result_ot;
+  ASSERT_TRUE(sampling_result.trace_state->Get("ot", result_ot));
+  ASSERT_EQ(ot_value, result_ot);
+}
+
+TEST(ProbabilitySampler, OversizedValueErasesStaleThreshold)
+{
+  // 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};
+
+  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);
+
+  auto sampling_result = SampleWithContext(s1, context, TraceIdWithRandomness(0xffffffffffffff));
+
+  ASSERT_EQ(Decision::RECORD_AND_SAMPLE, sampling_result.decision);
+  ASSERT_NE(nullptr, sampling_result.trace_state);
+
+  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:"));
 }
 
 TEST(ProbabilitySampler, IgnoresParentSampledFlag)