From 30f8a0478a3f7380a8f0e2e6f366c1f271cbeb46 Mon Sep 17 00:00:00 2001
From: Wangchong Zhou <wangchong@konghq.com>
Date: Fri, 5 Jul 2024 18:21:51 +0800
Subject: [PATCH] feat(pdk): add service.set_target_retry_callback

---
 kong/pdk/service.lua                          | 22 +++++++++++++
 kong/runloop/balancer/init.lua                | 10 ++++++
 t/01-pdk/06-service-request/00-phase_checks.t |  4 +--
 t/01-pdk/09-service/00-phase_checks.t         | 31 ++++++++++++++++---
 4 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/kong/pdk/service.lua b/kong/pdk/service.lua
index 59b0f9b62036..182c771753dd 100644
--- a/kong/pdk/service.lua
+++ b/kong/pdk/service.lua
@@ -107,6 +107,28 @@ local function new()
     ctx.balancer_data.port = port
   end
 
+
+  -- Sets the retry callback function when the target set by `service.set_target`
+  -- failed to connect. The callback function will be called with no argument and
+  -- must return `host`, `port` and `err` if any.
+  --
+  --
+  -- @function kong.service.set_target_retry_callback
+  -- @phases access
+  -- @tparam function retry_callback
+  -- @usage
+  -- kong.service.set_target_retry_callback(function() return "service.local", 443 end)
+  function service.set_target_retry_callback(retry_callback)
+    check_phase(PHASES.access)
+
+    if type(retry_callback) ~= "function" then
+      error("retry_callback must be a function", 2)
+    end
+
+    ngx.ctx.balancer_data.retry_callback = retry_callback
+  end
+
+
   ---
   -- Sets the retries count for the current request. This will override the
   -- default retries count set in the Upstream entity.
diff --git a/kong/runloop/balancer/init.lua b/kong/runloop/balancer/init.lua
index 550c1055d84e..51ad4872b5af 100644
--- a/kong/runloop/balancer/init.lua
+++ b/kong/runloop/balancer/init.lua
@@ -359,6 +359,16 @@ local function execute(balancer_data, ctx)
     balancer_data.balancer_handle = handle
 
   else
+    -- Note: balancer_data.retry_callback is only set by PDK once in access phase
+    -- if kong.service.set_target_retry_callback is called
+    if balancer_data.try_count ~= 0 and balancer_data.retry_callback then
+      local pok, perr, err = pcall(balancer_data.retry_callback)
+      if not pok or not perr then
+        log(ERR, "retry handler failed: ", err or perr)
+        return nil, "failure to get a peer from retry handler", 503
+      end
+    end
+
     -- have to do a regular DNS lookup
     local try_list
     local hstate = run_hook("balancer:to_ip:pre", balancer_data.host)
diff --git a/t/01-pdk/06-service-request/00-phase_checks.t b/t/01-pdk/06-service-request/00-phase_checks.t
index 80a57cbce4bc..d459338a9b21 100644
--- a/t/01-pdk/06-service-request/00-phase_checks.t
+++ b/t/01-pdk/06-service-request/00-phase_checks.t
@@ -47,7 +47,7 @@ qq{
                 args          = { "http" },
                 init_worker   = false,
                 certificate   = "pending",
-                rewrite       = "forced false",
+                rewrite       = true,
                 access        = true,
                 response      = "forced false",
                 header_filter = "forced false",
@@ -71,7 +71,7 @@ qq{
                 args          = { "/" },
                 init_worker   = false,
                 certificate   = "pending",
-                rewrite       = "forced false",
+                rewrite       = true,
                 access        = true,
                 response      = "forced false",
                 header_filter = "forced false",
diff --git a/t/01-pdk/09-service/00-phase_checks.t b/t/01-pdk/09-service/00-phase_checks.t
index 87869a826313..c19d9a824b70 100644
--- a/t/01-pdk/09-service/00-phase_checks.t
+++ b/t/01-pdk/09-service/00-phase_checks.t
@@ -88,6 +88,18 @@ qq{
                 body_filter   = "forced false",
                 log           = "forced false",
                 admin_api     = "forced false",
+            }, {
+                method        = "set_target_retry_callback",
+                args          = { function() end },
+                init_worker   = "forced false",
+                certificate   = "pending",
+                rewrite       = "forced false",
+                access        = true,
+                response      = "forced false",
+                header_filter = "forced false",
+                body_filter   = "forced false",
+                log           = "forced false",
+                admin_api     = "forced false",
             }, {
                 method        = "set_timeouts",
                 args          = { 1, 2, 3},
@@ -251,8 +263,7 @@ qq{
                 log           = "pending",
                 admin_api     = "forced false",
                 preread       = "pending",
-            },
-            {
+            }, {
                 method        = "set_retries",
                 args          = { 3, },
                 init_worker   = "forced false",
@@ -265,6 +276,19 @@ qq{
                 log           = "pending",
                 admin_api     = "forced false",
                 preread       = "pending",
+            }, {
+                method        = "set_target_retry_callback",
+                args          = { function() end },
+                init_worker   = "forced false",
+                certificate   = "pending",
+                rewrite       = "forced false",
+                access        = true,
+                response      = "forced false",
+                header_filter = "forced false",
+                body_filter   = "forced false",
+                log           = "pending",
+                admin_api     = "forced false",
+                preread       = "pending",
             }, {
                 method        = "set_timeouts",
                 args          = { 1, 2, 3},
@@ -278,8 +302,7 @@ qq{
                 log           = "pending",
                 admin_api     = "forced false",
                 preread       = "pending",
-            }, 
-            {
+            }, {
                 method        = "set_tls_cert_key",
                 args          = { chain, key, },
                 init_worker   = false,