From b5ba45387247a8d9c1b48bf654270db3818a84e3 Mon Sep 17 00:00:00 2001
From: terry-xiaoyu <506895667@qq.com>
Date: Thu, 9 Apr 2020 19:19:31 +0800
Subject: [PATCH] Add add_reconnect_callback/2

---
 src/ecpool.erl        |  8 +++++++-
 src/ecpool_worker.erl | 12 ++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/ecpool.erl b/src/ecpool.erl
index 2d73ad002..58b587e91 100644
--- a/src/ecpool.erl
+++ b/src/ecpool.erl
@@ -28,7 +28,7 @@
         , workers/1
         ]).
 
--export([set_reconnect_callback/2]).
+-export([set_reconnect_callback/2, add_reconnect_callback/2]).
 
 -export_type([ pool_name/0
              , pool_type/0
@@ -84,6 +84,12 @@ set_reconnect_callback(Pool, Callback) ->
      || {_WorkerName, Worker} <- ecpool:workers(Pool)],
     ok.
 
+-spec(add_reconnect_callback(atom(), reconn_callback()) -> ok).
+add_reconnect_callback(Pool, Callback) ->
+    [ecpool_worker:add_reconnect_callback(Worker, Callback)
+     || {_WorkerName, Worker} <- ecpool:workers(Pool)],
+    ok.
+
 %% @doc Call the fun with client/connection
 -spec(with_client(atom(), fun((Client :: pid()) -> any())) -> no_return()).
 with_client(Pool, Fun) when is_atom(Pool) ->
diff --git a/src/ecpool_worker.erl b/src/ecpool_worker.erl
index 506f0b1df..fe048e9d1 100644
--- a/src/ecpool_worker.erl
+++ b/src/ecpool_worker.erl
@@ -24,6 +24,7 @@
 -export([ client/1
         , is_connected/1
         , set_reconnect_callback/2
+        , add_reconnect_callback/2
         ]).
 
 %% gen_server Function Exports
@@ -90,6 +91,10 @@ is_connected(Pid) ->
 set_reconnect_callback(Pid, OnReconnect) ->
     gen_server:cast(Pid, {set_reconn_callbk, OnReconnect}).
 
+-spec(add_reconnect_callback(pid(), ecpool:reconn_callback()) -> ok).
+add_reconnect_callback(Pid, OnReconnect) ->
+    gen_server:cast(Pid, {add_reconn_callbk, OnReconnect}).
+
 %%--------------------------------------------------------------------
 %% gen_server callbacks
 %%--------------------------------------------------------------------
@@ -130,6 +135,9 @@ handle_call(Req, _From, State) ->
 handle_cast({set_reconn_callbk, OnReconnect}, State) ->
     {noreply, State#state{on_reconnect = OnReconnect}};
 
+handle_cast({add_reconn_callbk, OnReconnect}, State = #state{on_reconnect = OnReconnectList}) ->
+    {noreply, State#state{on_reconnect = [OnReconnect | OnReconnectList]}};
+
 handle_cast(_Msg, State) ->
     {noreply, State}.
 
@@ -195,6 +203,10 @@ reconnect(Secs, State = #state{client = Client, on_disconnect = Disconnect, supe
 
 handle_reconnect(_, undefined) ->
     ok;
+handle_reconnect(undefined, _) ->
+    ok;
+handle_reconnect(Client, OnReconnectList) when is_list(OnReconnectList) ->
+    [OnReconnect(Client) || OnReconnect <- OnReconnectList];
 handle_reconnect(Client, OnReconnect) ->
     OnReconnect(Client).