From 16a2d4dfa97a8b10cf37936aed8c0173b374d4a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=BD=97=E6=B3=BD=E8=BD=A9?= <spacewanderlzx@gmail.com>
Date: Mon, 5 Oct 2020 23:19:40 +0800
Subject: [PATCH] feat: also call picker.after_balance before retry (#2350)

---
 apisix/balancer.lua      | 4 ++++
 apisix/balancer/ewma.lua | 7 ++++++-
 apisix/init.lua          | 2 +-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/apisix/balancer.lua b/apisix/balancer.lua
index cfe53b95..f2eec902 100644
--- a/apisix/balancer.lua
+++ b/apisix/balancer.lua
@@ -172,6 +172,10 @@ local function pick_server(route, ctx)
         return nil, "no valid upstream node"
     end
 
+    if ctx.server_picker and ctx.server_picker.after_balance then
+        ctx.server_picker.after_balance(ctx, true)
+    end
+
     if up_conf.timeout then
         local timeout = up_conf.timeout
         local ok, err = set_timeouts(timeout.connect, timeout.send,
diff --git a/apisix/balancer/ewma.lua b/apisix/balancer/ewma.lua
index 5198c282..bdc7722e 100644
--- a/apisix/balancer/ewma.lua
+++ b/apisix/balancer/ewma.lua
@@ -146,7 +146,12 @@ local function _ewma_find(ctx, up_nodes)
 end
 
 
-local function _ewma_after_balance(ctx)
+local function _ewma_after_balance(ctx, before_retry)
+    if before_retry then
+        -- don't count tries which fail to complete
+        return nil
+    end
+
     local response_time = tonumber(ctx.var.upstream_response_time) or 0
     local connect_time = tonumber(ctx.var.upstream_connect_time) or 0
     local rtt = connect_time + response_time
diff --git a/apisix/init.lua b/apisix/init.lua
index 780757f2..cc7902b5 100644
--- a/apisix/init.lua
+++ b/apisix/init.lua
@@ -672,7 +672,7 @@ function _M.http_log_phase()
     healcheck_passive(api_ctx)
 
     if api_ctx.server_picker and api_ctx.server_picker.after_balance then
-        api_ctx.server_picker.after_balance(api_ctx)
+        api_ctx.server_picker.after_balance(api_ctx, false)
     end
 
     if api_ctx.uri_parse_param then