From 1477424e32c0a273a4bbe01cb597c36cce397fd8 Mon Sep 17 00:00:00 2001
From: Shawn <506895667@qq.com>
Date: Sat, 22 Jun 2024 14:22:18 +0800
Subject: [PATCH] fix: fill up the global client tab when relup

---
 src/ecpool.appup.src   | 12 +++++++++---
 src/ecpool.erl         |  9 +++++++--
 src/ecpool_monitor.erl | 43 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/src/ecpool.appup.src b/src/ecpool.appup.src
index abdda4fc..5a2a1f25 100644
--- a/src/ecpool.appup.src
+++ b/src/ecpool.appup.src
@@ -5,6 +5,7 @@
       %% NOTE: MUST start ecpool_monitor before any load_module instructions
       {add_module, ecpool_monitor},
       {apply, {ecpool_monitor, ensure_monitor_started, []}},
+      {apply, {ecpool_monitor, update_clients_global, []}},
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool_sup, brutal_purge, soft_purge, []},
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
@@ -13,6 +14,7 @@
     {<<"0\\.5\\.[0-2]">>, [
       {add_module, ecpool_monitor},
       {apply, {ecpool_monitor, ensure_monitor_started, []}},
+      {apply, {ecpool_monitor, update_clients_global, []}},
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool_sup, brutal_purge, soft_purge, []},
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
@@ -22,6 +24,7 @@
     {"0.4.2", [
       {add_module, ecpool_monitor},
       {apply, {ecpool_monitor, ensure_monitor_started, []}},
+      {apply, {ecpool_monitor, update_clients_global, []}},
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool, brutal_purge, soft_purge, []},
@@ -34,21 +37,24 @@
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool_sup, brutal_purge, soft_purge, []},
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
-      {load_module, ecpool, brutal_purge, soft_purge, []}
+      {load_module, ecpool, brutal_purge, soft_purge, []},
+      {apply, {ecpool_monitor, ensure_monitor_stopped, []}}
     ]},
     {<<"0\\.5\\.[0-2]">>, [
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool_sup, brutal_purge, soft_purge, []},
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
       {load_module, ecpool, brutal_purge, soft_purge, []},
-      {load_module, ecpool_worker_sup, brutal_purge, soft_purge, []}
+      {load_module, ecpool_worker_sup, brutal_purge, soft_purge, []},
+      {apply, {ecpool_monitor, ensure_monitor_stopped, []}}
     ]},
     {"0.4.2", [
       {load_module, ecpool_worker, brutal_purge, soft_purge, []},
       {load_module, ecpool_pool, brutal_purge, soft_purge, []},
       {load_module, ecpool, brutal_purge, soft_purge, []},
       {load_module, ecpool_worker_sup, brutal_purge, soft_purge, []},
-      {load_module, ecpool_sup, brutal_purge, soft_purge, []}
+      {load_module, ecpool_sup, brutal_purge, soft_purge, []},
+      {apply, {ecpool_monitor, ensure_monitor_stopped, []}}
     ]}
   ]
 }.
diff --git a/src/ecpool.erl b/src/ecpool.erl
index 28466c4a..fb0409a6 100644
--- a/src/ecpool.erl
+++ b/src/ecpool.erl
@@ -155,8 +155,13 @@ with_worker(false, Action, _Mode) when ?IS_ACTION(Action) ->
     {error, ecpool_empty};
 with_worker(Worker, Action, no_handover) when ?IS_ACTION(Action) ->
     case ecpool_monitor:get_client_global(Worker) of
-        {ok, Client} -> exec(Action, Client);
-        {error, Reason} -> {error, Reason}
+        {ok, Client} ->
+            exec(Action, Client);
+        {error, _} ->
+            case ecpool_worker:client(Worker) of
+                {ok, Client} -> exec(Action, Client);
+                {error, _} = Err -> Err
+            end
     end;
 with_worker(Worker, Action, handover) when ?IS_ACTION(Action) ->
     ecpool_worker:exec(Worker, Action, infinity);
diff --git a/src/ecpool_monitor.erl b/src/ecpool_monitor.erl
index 48985f85..0fbf36af 100644
--- a/src/ecpool_monitor.erl
+++ b/src/ecpool_monitor.erl
@@ -19,8 +19,19 @@
 -behaviour(gen_server).
 
 %% API
--export([start_link/0, ensure_monitor_started/0, monitor_spec/0]).
--export([reg_worker/0, put_client_global/1, get_client_global/1]).
+-export([ start_link/0
+        , ensure_monitor_started/0
+        , ensure_monitor_stopped/0
+        , monitor_spec/0
+        ]).
+
+-export([ update_clients_global/0
+        , reg_worker/0
+        , get_all_global_clients/0
+        , put_client_global/1
+        , put_client_global/2
+        , get_client_global/1
+        ]).
 
 %% gen_server callbacks
 -export([init/1,
@@ -28,7 +39,8 @@
          handle_cast/2,
          handle_info/2,
          terminate/2,
-         code_change/3]).
+         code_change/3
+        ]).
 
 -define(DISCOVERY_TAB, ecpool_global_client_discovery).
 
@@ -42,6 +54,13 @@ ensure_monitor_started() ->
         {error, _} = Error -> Error
     end.
 
+ensure_monitor_stopped() ->
+    case supervisor:terminate_child(ecpool_sup, ecpool_monitor) of
+        ok -> supervisor:delete_child(ecpool_sup, ecpool_monitor);
+        {error, not_found} -> ok;
+        Error -> Error
+    end.
+
 monitor_spec() ->
     #{
         id => ecpool_monitor,
@@ -52,8 +71,24 @@ monitor_spec() ->
         modules => [ecpool_monitor]
     }.
 
+update_clients_global() ->
+    lists:foreach(fun({PoolName, _}) ->
+        lists:foreach(fun({_, WrokerPid}) ->
+                case ecpool_worker:client(WrokerPid) of
+                    {ok, Client} -> put_client_global(WrokerPid, Client);
+                    _ -> ok
+                end
+            end, ecpool:workers(PoolName))
+    end, ecpool_sup:pools()).
+
+get_all_global_clients() ->
+    ets:tab2list(?DISCOVERY_TAB).
+
 put_client_global(Client) ->
-    ets:insert(?DISCOVERY_TAB, {self(), Client}),
+    put_client_global(self(), Client).
+
+put_client_global(WrokerPid, Client) ->
+    ets:insert(?DISCOVERY_TAB, {WrokerPid, Client}),
     ok.
 
 get_client_global(WorkerPid) ->