fix(tracing): ensure sampling rate applies to whole trace (#11135)

* fix(tracing): ensure sampling rate applies to whole trace

This commit fixes a bug where the sampling rate was applied to
individual spans, so configuring values of sampling_rate < 1 would
sometimes result in split traces being reported.

The fix ensures that the sampled flag that is calculated by the sampler
for the root span is propagated to all children within the same trace.

* fix(tracing): suggestions + propagation test

CHANGELOG.md

@@ -71,6 +71,8 @@
   [#11066](https://github.com/Kong/kong/pull/11066)
 - Remove kong branding from kong HTML error template.
   [#11150](https://github.com/Kong/kong/pull/11150)
+- Fix a bug that caused sampling rate to be applied to individual spans producing split traces.
+  [#11135](https://github.com/Kong/kong/pull/11135)
 
 #### Admin API
 

kong/pdk/tracing.lua

@@ -195,17 +195,13 @@ local function create_span(tracer, options)
     sampled = tracer and tracer.sampler(trace_id)
   end
 
-  if not sampled then
-    return noop_span
-  end
-
   span.parent_id = span.parent and span.parent.span_id
       or options.parent_id
   span.tracer = span.tracer or tracer
   span.span_id = generate_span_id()
   span.trace_id = trace_id
   span.kind = options.span_kind or SPAN_KIND.INTERNAL
-  span.should_sample = true
+  span.should_sample = sampled
 
   setmetatable(span, span_mt)
   return span
@@ -275,8 +271,8 @@ end
 -- local time = ngx.now()
 -- span:finish(time * 100000000)
 function span_mt:finish(end_time_ns)
-  if self.end_time_ns ~= nil then
-    -- span is finished, and processed already
+  if self.end_time_ns ~= nil or not self.should_sample then
+    -- span is finished, and already processed or not sampled
     return
   end
 

spec/02-integration/14-tracing/02-propagation_spec.lua

@@ -160,7 +160,7 @@ for _, strategy in helpers.each_strategy() do
         assert.is_same(0, #spans, res)
 
         local traceparent = assert(body.headers.traceparent)
-        assert.equals("00-" .. trace_id .. "-" .. span_id .. "-00", traceparent)
+        assert.matches("00%-" .. trace_id .. "%-%x+%-00", traceparent)
       end)
     end)
   end)

spec/02-integration/14-tracing/03-tracer-pdk_spec.lua

@@ -0,0 +1,81 @@
+local helpers = require "spec.helpers"
+local cjson   = require "cjson"
+local join    = require "pl.stringx".join
+
+
+local TCP_PORT = helpers.get_available_port()
+local tcp_trace_plugin_name = "tcp-trace-exporter"
+
+
+for _, strategy in helpers.each_strategy() do
+  local proxy_client
+
+  describe("tracer pdk spec #" .. strategy, function()
+
+    local function setup_instrumentations(types, custom_spans, sampling_rate)
+      local bp, _ = assert(helpers.get_db_utils(strategy, {
+        "services",
+        "routes",
+        "plugins",
+      }, { tcp_trace_plugin_name }))
+
+      local http_srv = assert(bp.services:insert {
+        name = "mock-service",
+        host = helpers.mock_upstream_host,
+        port = helpers.mock_upstream_port,
+      })
+
+      bp.routes:insert({ service = http_srv,
+                         protocols = { "http" },
+                         paths = { "/" }})
+
+      bp.plugins:insert({
+        name = tcp_trace_plugin_name,
+        config = {
+          host = "127.0.0.1",
+          port = TCP_PORT,
+          custom_spans = custom_spans or false,
+        }
+      })
+
+      assert(helpers.start_kong {
+        database = strategy,
+        nginx_conf = "spec/fixtures/custom_nginx.template",
+        plugins = "tcp-trace-exporter",
+        tracing_instrumentations = types,
+        tracing_sampling_rate = sampling_rate or 1,
+      })
+
+      proxy_client = helpers.proxy_client()
+    end
+
+    describe("sampling rate", function ()
+      local instrumentations = { "request", "router", "balancer" }
+      lazy_setup(function()
+        setup_instrumentations(join(",", instrumentations), false, 0.5)
+      end)
+
+      lazy_teardown(function()
+        helpers.stop_kong()
+      end)
+
+      it("results in either all or none of the spans in a trace to be sampled", function ()
+        for _ = 1, 100 do
+          local thread = helpers.tcp_server(TCP_PORT)
+          local r = assert(proxy_client:send {
+            method  = "GET",
+            path    = "/",
+          })
+          assert.res_status(200, r)
+
+          local ok, res = thread:join()
+          assert.True(ok)
+          assert.is_string(res)
+
+          local spans = cjson.decode(res)
+          assert.True(#spans == 0 or #spans == #instrumentations)
+        end
+      end)
+    end)
+  end)
+end