fix(tracing/core): sampled flags and propagation

kong/pdk/tracing.lua

@@ -156,9 +156,17 @@ local function create_span(tracer, options)
   local trace_id = span.parent and span.parent.trace_id
       or options.trace_id
       or generate_trace_id()
-  local sampled = span.parent and span.parent.should_sample
-      or options.should_sample
-      or tracer and tracer.sampler(trace_id)
+
+  local sampled
+  if span.parent and span.parent.should_sample ~= nil then
+    sampled = span.parent.should_sample
+
+  elseif options.should_sample ~= nil then
+    sampled = options.should_sample
+
+  else
+    sampled = tracer and tracer.sampler(trace_id)
+  end
 
   if not sampled then
     return noop_span
@@ -170,16 +178,17 @@ local function create_span(tracer, options)
   span.span_id = generate_span_id()
   span.trace_id = trace_id
   span.kind = options.span_kind or SPAN_KIND.INTERNAL
-  -- indicates whether the span should be reported
-  span.should_sample = span.parent and span.parent.should_sample
-      or options.should_sample
-      or sampled
+  span.should_sample = true
 
   setmetatable(span, span_mt)
   return span
 end
 
 local function link_span(tracer, span, name, options)
+  if not span.should_sample then
+    kong.log.debug("skipping non-sampled span")
+    return
+  end
   if tracer and type(tracer) ~= "table" then
     error("invalid tracer", 2)
   end
@@ -386,6 +395,7 @@ noop_tracer.link_span = NOOP
 noop_tracer.active_span = NOOP
 noop_tracer.set_active_span = NOOP
 noop_tracer.process_span = NOOP
+noop_tracer.set_should_sample = NOOP
 
 --- New Tracer
 local function new_tracer(name, options)
@@ -482,12 +492,29 @@ local function new_tracer(name, options)
     end
 
     for _, span in ipairs(ctx.KONG_SPANS) do
-      if span.tracer.name == self.name then
+      if span.tracer and span.tracer.name == self.name then
         processor(span, ...)
       end
     end
   end
 
+  --- Update the value of should_sample for all spans
+  --
+  -- @function span:set_should_sample
+  -- @tparam bool should_sample value for the sample parameter
+  function self:set_should_sample(should_sample)
+    local ctx = ngx.ctx
+    if not ctx.KONG_SPANS then
+      return
+    end
+
+    for _, span in ipairs(ctx.KONG_SPANS) do
+      if span.is_recording ~= false then
+        span.should_sample = should_sample
+      end
+    end
+  end
+
   tracer_memo[name] = setmetatable(self, tracer_mt)
   return tracer_memo[name]
 end

kong/tracing/propagation.lua

@@ -427,6 +427,11 @@ end
 
 
 local function set(conf_header_type, found_header_type, proxy_span, conf_default_header_type)
+  if proxy_span.is_recording == false then
+    kong.log.debug("skipping propagation of noop span")
+    return
+  end
+
   local set_header = kong.service.request.set_header
 
   -- If conf_header_type is set to `preserve`, found_header_type is used over default_header_type;

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

@@ -1,47 +1,59 @@
 local helpers = require "spec.helpers"
 local cjson = require "cjson"
+local utils = require "kong.tools.utils"
+local to_hex = require("resty.string").to_hex
+
+local rand_bytes = utils.get_rand_bytes
 
 local TCP_PORT = 35001
+
+local function gen_id(len)
+  return to_hex(rand_bytes(len))
+end
+
 for _, strategy in helpers.each_strategy() do
   local proxy_client
 
   describe("tracing propagation spec #" .. strategy, function()
-    describe("spans hierarchy", function ()
+    lazy_setup(function()
+      local bp, _ = assert(helpers.get_db_utils(strategy, {
+        "routes",
+        "plugins",
+      }, { "tcp-trace-exporter", "trace-propagator" }))
 
-      lazy_setup(function()
-        local bp, _ = assert(helpers.get_db_utils(strategy, {
-          "routes",
-          "plugins",
-        }, { "tcp-trace-exporter", "trace-propagator" }))
-
-        bp.routes:insert({
-          hosts = { "propagate.test" },
-        })
+      bp.routes:insert({
+        hosts = { "propagate.test" },
+      })
 
-        bp.plugins:insert({
-          name = "tcp-trace-exporter",
-          config = {
-            host = "127.0.0.1",
-            port = TCP_PORT,
-            custom_spans = false,
-          }
-        })
+      bp.plugins:insert({
+        name = "tcp-trace-exporter",
+        config = {
+          host = "127.0.0.1",
+          port = TCP_PORT,
+          custom_spans = false,
+        }
+      })
 
-        bp.plugins:insert({
-          name = "trace-propagator"
-        })
+      bp.plugins:insert({
+        name = "trace-propagator"
+      })
+    end)
 
+    describe("spans hierarchy", function ()
+      lazy_setup(function()
         assert(helpers.start_kong {
           database = strategy,
           nginx_conf = "spec/fixtures/custom_nginx.template",
           plugins = "tcp-trace-exporter,trace-propagator",
           tracing_instrumentations = "balancer",
         })
-
         proxy_client = helpers.proxy_client()
       end)
 
       lazy_teardown(function()
+        if proxy_client then
+          proxy_client:close()
+        end
         helpers.stop_kong()
       end)
 
@@ -75,5 +87,79 @@ for _, strategy in helpers.each_strategy() do
         assert.equals("00-" .. trace_id .. "-" .. span_id .. "-01", traceparent)
       end)
     end)
+
+    describe("parsing incoming headers", function ()
+      local trace_id = gen_id(16)
+      local span_id = gen_id(8)
+
+      lazy_setup(function()
+        assert(helpers.start_kong {
+          database = strategy,
+          nginx_conf = "spec/fixtures/custom_nginx.template",
+          plugins = "tcp-trace-exporter,trace-propagator",
+          tracing_instrumentations = "request,router,balancer,plugin_access,plugin_header_filter",
+        })
+        proxy_client = helpers.proxy_client()
+      end)
+
+      lazy_teardown(function()
+        if proxy_client then
+          proxy_client:close()
+        end
+        helpers.stop_kong()
+      end)
+
+      it("enables sampling when incoming header has sampled enabled", function ()
+        local thread = helpers.tcp_server(TCP_PORT)
+        local r = assert(proxy_client:send {
+          method  = "GET",
+          path = "/request",
+          headers = {
+            ["Host"] = "propagate.test",
+            traceparent = string.format("00-%s-%s-01", trace_id, span_id),
+          }
+        })
+        assert.res_status(200, r)
+        local body = r:read_body()
+        body = assert(body and cjson.decode(body))
+
+        local ok, res = thread:join()
+        assert.True(ok)
+        assert.is_string(res)
+
+        -- all spans are returned
+        local spans = cjson.decode(res)
+        assert.is_same(6, #spans, res)
+
+        local traceparent = assert(body.headers.traceparent)
+        assert.matches("00%-" .. trace_id .. "%-%x+%-01", traceparent)
+      end)
+
+      it("disables sampling when incoming header has sampled disabled", function ()
+        local thread = helpers.tcp_server(TCP_PORT)
+        local r = assert(proxy_client:send {
+          method  = "GET",
+          path = "/request",
+          headers = {
+            ["Host"] = "propagate.test",
+            traceparent = string.format("00-%s-%s-00", trace_id, span_id),
+          }
+        })
+        assert.res_status(200, r)
+        local body = r:read_body()
+        body = assert(body and cjson.decode(body))
+
+        local ok, res = thread:join()
+        assert.True(ok)
+        assert.is_string(res)
+
+        -- no spans are returned
+        local spans = cjson.decode(res)
+        assert.is_same(0, #spans, res)
+
+        local traceparent = assert(body.headers.traceparent)
+        assert.equals("00-" .. trace_id .. "-" .. span_id .. "-00", traceparent)
+      end)
+    end)
   end)
 end

spec/fixtures/custom_plugins/kong/plugins/tcp-trace-exporter/handler.lua

@@ -104,6 +104,9 @@ function _M:log(config)
 
   local spans = {}
   local process_span = function (span)
+    if span.should_sample == false then
+      return
+    end
     local s = table.clone(span)
     s.tracer = nil
     s.parent = nil

spec/fixtures/custom_plugins/kong/plugins/trace-propagator/handler.lua

@@ -16,7 +16,11 @@ function _M:access(conf)
   local tracer = kong.tracing.new("trace-propagator")
   local root_span = tracer.start_span("root")
 
-  local header_type, trace_id, span_id, parent_id = propagation_parse(headers)
+  local header_type, trace_id, span_id, parent_id, should_sample = propagation_parse(headers)
+
+  if should_sample == false then
+    tracer:set_should_sample(should_sample)
+  end
 
   if trace_id then
     root_span.trace_id = trace_id