From 65ef8ca3dfb6ca89628b4a20e2cd37ba0d624f6b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=BD=97=E6=B3=BD=E8=BD=A9?= <spacewanderlzx@gmail.com>
Date: Sat, 6 Feb 2021 13:18:28 +0800
Subject: [PATCH] fix(zipkin): don't cache the per-req sample ratio (#3522)

Signed-off-by: spacewander <spacewanderlzx@gmail.com>
---
 apisix/plugins/zipkin.lua                | 42 +++++++++++------------
 apisix/plugins/zipkin/random_sampler.lua | 10 +++---
 t/plugin/zipkin.t                        | 43 ++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 27 deletions(-)

diff --git a/apisix/plugins/zipkin.lua b/apisix/plugins/zipkin.lua
index fddcd6c3cfc2..44f241f35e1e 100644
--- a/apisix/plugins/zipkin.lua
+++ b/apisix/plugins/zipkin.lua
@@ -64,24 +64,6 @@ end
 
 
 local function create_tracer(conf,ctx)
-
-    local headers = core.request.headers(ctx)
-
-    -- X-B3-Sampled: if the client decided to sample this request, we do too.
-    local sample = headers["x-b3-sampled"]
-    if sample == "1" or sample == "true" then
-        conf.sample_ratio = 1
-    elseif sample == "0" or sample == "false" then
-        conf.sample_ratio = 0
-    end
-
-    -- X-B3-Flags: if it equals '1' then it overrides sampling policy
-    -- We still want to warn on invalid sample header, so do this after the above
-    local debug = headers["x-b3-flags"]
-    if debug == "1" then
-        conf.sample_ratio = 1
-    end
-
     conf.route_id = ctx.route_id
     local reporter = new_reporter(conf)
     reporter:init_processor()
@@ -100,10 +82,28 @@ function _M.rewrite(plugin_conf, ctx)
         conf.server_addr = ctx.var["server_addr"]
     end
 
-    local tracer = core.lrucache.plugin_ctx(lrucache, ctx, conf.server_addr,
-                            create_tracer, conf, ctx)
+    local tracer = core.lrucache.plugin_ctx(lrucache, ctx, conf.server_addr .. conf.server_port,
+                                            create_tracer, conf, ctx)
+
+    local headers = core.request.headers(ctx)
+    local per_req_sample_ratio
+
+    -- X-B3-Sampled: if the client decided to sample this request, we do too.
+    local sample = headers["x-b3-sampled"]
+    if sample == "1" or sample == "true" then
+        per_req_sample_ratio = 1
+    elseif sample == "0" or sample == "false" then
+        per_req_sample_ratio = 0
+    end
+
+    -- X-B3-Flags: if it equals '1' then it overrides sampling policy
+    -- We still want to warn on invalid sample header, so do this after the above
+    local debug = headers["x-b3-flags"]
+    if debug == "1" then
+        per_req_sample_ratio = 1
+    end
 
-    ctx.opentracing_sample = tracer.sampler:sample()
+    ctx.opentracing_sample = tracer.sampler:sample(per_req_sample_ratio or conf.sample_ratio)
     if not ctx.opentracing_sample then
         core.request.set_header("x-b3-sampled", "0")
         return
diff --git a/apisix/plugins/zipkin/random_sampler.lua b/apisix/plugins/zipkin/random_sampler.lua
index f2ee5b49a550..cfed3f914b07 100644
--- a/apisix/plugins/zipkin/random_sampler.lua
+++ b/apisix/plugins/zipkin/random_sampler.lua
@@ -24,16 +24,14 @@ local _M = {}
 local mt = { __index = _M }
 
 function _M.new(conf)
-    local sample_ratio = conf.sample_ratio
-    assert(type(sample_ratio) == "number" and
-        sample_ratio >= 0 and sample_ratio <= 1, "invalid sample_ratio")
     return setmetatable({
-        sample_ratio = sample_ratio
 	}, mt)
 end
 
-function _M.sample(self)
-    return math.random() < self.sample_ratio
+function _M.sample(self, sample_ratio)
+    assert(type(sample_ratio) == "number" and
+           sample_ratio >= 0 and sample_ratio <= 1, "invalid sample_ratio")
+    return math.random() < sample_ratio
 end
 
 
diff --git a/t/plugin/zipkin.t b/t/plugin/zipkin.t
index 720412b27d1b..89f6725f356f 100644
--- a/t/plugin/zipkin.t
+++ b/t/plugin/zipkin.t
@@ -559,3 +559,46 @@ GET /echo
 x-b3-sampled: true
 --- response_headers
 x-b3-sampled: 1
+
+
+
+=== TEST 23: don't cache the per-req sample ratio
+--- config
+    location /t {
+        content_by_lua_block {
+            local http = require "resty.http"
+            local httpc = http.new()
+            local uri = "http://127.0.0.1:" .. ngx.var.server_port
+                        .. "/echo"
+            -- force to trace
+            local res, err = httpc:request_uri(uri, {
+                method = "GET",
+                headers = {
+                    ['x-b3-sampled'] = 1
+                }
+            })
+            if not res then
+                ngx.say(err)
+                return
+            end
+            ngx.say(res.headers['x-b3-sampled'])
+
+            -- force not to trace
+            local res, err = httpc:request_uri(uri, {
+                method = "GET",
+                headers = {
+                    ['x-b3-sampled'] = 0
+                }
+            })
+            if not res then
+                ngx.say(err)
+                return
+            end
+            ngx.say(res.headers['x-b3-sampled'])
+        }
+    }
+--- request
+GET /t
+--- response_body
+1
+0