From 0d8c137cb558d088c69fa2aceb0a5e01b5578ddd Mon Sep 17 00:00:00 2001
From: JianBo He <heeejianbo@gmail.com>
Date: Thu, 25 Feb 2021 19:55:54 +0800
Subject: [PATCH] refactor(exhook): load only the necessary hooks

---
 apps/emqx_exhook/src/emqx_exhook_app.erl     | 19 ++++++----
 apps/emqx_exhook/src/emqx_exhook_handler.erl | 22 ++++++++++-
 apps/emqx_exhook/src/emqx_exhook_server.erl  | 39 ++++++++++++++++++--
 3 files changed, 68 insertions(+), 12 deletions(-)

diff --git a/apps/emqx_exhook/src/emqx_exhook_app.erl b/apps/emqx_exhook/src/emqx_exhook_app.erl
index b008c251a..45474d925 100644
--- a/apps/emqx_exhook/src/emqx_exhook_app.erl
+++ b/apps/emqx_exhook/src/emqx_exhook_app.erl
@@ -22,6 +22,8 @@
 
 -emqx_plugin(extension).
 
+-define(REGISTRAY, emqx_exhook_registray).
+
 -export([ start/2
         , stop/1
         , prep_stop/1
@@ -30,8 +32,8 @@
 %% Internal export
 -export([ load_server/2
         , unload_server/1
-        , load_exhooks/0
         , unload_exhooks/0
+        , init_hook_registray/0
         ]).
 
 %%--------------------------------------------------------------------
@@ -41,12 +43,12 @@
 start(_StartType, _StartArgs) ->
     {ok, Sup} = emqx_exhook_sup:start_link(),
 
+    %% Collect all available hooks
+    _ = init_hook_registray(),
+
     %% Load all dirvers
     load_all_servers(),
 
-    %% Register all hooks
-    _ = load_exhooks(),
-
     %% Register CLI
     emqx_ctl:register_command(exhook, {emqx_exhook_cli, cli}, []),
     {ok, Sup}.
@@ -81,11 +83,14 @@ unload_server(Name) ->
 %%--------------------------------------------------------------------
 %% Exhooks
 
-load_exhooks() ->
-    [emqx:hook(Name, {M, F, A}) || {Name, {M, F, A}} <- search_exhooks()].
+init_hook_registray() ->
+    _ = ets:new(?REGISTRAY, [public, named_table]),
+    [ets:insert(?REGISTRAY, {Name, {M, F, A}, 0})
+     || {Name, {M, F, A}} <- search_exhooks()].
 
 unload_exhooks() ->
-    [emqx:unhook(Name, {M, F}) || {Name, {M, F, _A}} <- search_exhooks()].
+    [emqx:unhook(Name, {M, F}) ||
+     {Name, {M, F, _A}, _} <- ets:tab2list(?REGISTRAY)].
 
 search_exhooks() ->
     search_exhooks(ignore_lib_apps(application:loaded_applications())).
diff --git a/apps/emqx_exhook/src/emqx_exhook_handler.erl b/apps/emqx_exhook/src/emqx_exhook_handler.erl
index c565d5dee..980fb7786 100644
--- a/apps/emqx_exhook/src/emqx_exhook_handler.erl
+++ b/apps/emqx_exhook/src/emqx_exhook_handler.erl
@@ -307,8 +307,26 @@ stringfy(Term) when is_atom(Term) ->
 stringfy(Term) ->
     unicode:characters_to_binary((io_lib:format("~0p", [Term]))).
 
-hexstr(B) ->
-    iolist_to_binary([io_lib:format("~2.16.0B", [X]) || X <- binary_to_list(B)]).
+hexstr(<<>>) -> [];
+hexstr(<<H:4, L:4, B/binary>>) ->
+    iolist_to_binary([hexchar(H), hexchar(L)] ++ hexstr(B)).
+
+hexchar(0) -> $0;
+hexchar(1) -> $1;
+hexchar(2) -> $2;
+hexchar(3) -> $3;
+hexchar(4) -> $4;
+hexchar(5) -> $5;
+hexchar(6) -> $6;
+hexchar(7) -> $7;
+hexchar(8) -> $8;
+hexchar(9) -> $9;
+hexchar(10) -> $A;
+hexchar(11) -> $B;
+hexchar(12) -> $C;
+hexchar(13) -> $D;
+hexchar(14) -> $E;
+hexchar(15) -> $F.
 
 %%--------------------------------------------------------------------
 %% Acc funcs
diff --git a/apps/emqx_exhook/src/emqx_exhook_server.erl b/apps/emqx_exhook/src/emqx_exhook_server.erl
index 451983437..76a2e491d 100644
--- a/apps/emqx_exhook/src/emqx_exhook_server.erl
+++ b/apps/emqx_exhook/src/emqx_exhook_server.erl
@@ -20,6 +20,7 @@
 
 -logger_header("[ExHook Svr]").
 
+-define(REGISTRAY, emqx_exhook_registray).
 -define(PB_CLIENT_MOD, emqx_exhook_v_1_hook_provider_client).
 
 %% Load/Unload
@@ -83,13 +84,19 @@
 load(Name0, Opts0) ->
     Name = prefix(Name0),
     {SvrAddr, ClientOpts} = channel_opts(Opts0),
-    case emqx_exhook_sup:start_grpc_client_channel(Name, SvrAddr, ClientOpts) of
+    case emqx_exhook_sup:start_grpc_client_channel(
+           Name,
+           SvrAddr,
+           ClientOpts) of
         {ok, _ChannPoolPid} ->
             case do_init(Name) of
                 {ok, HookSpecs} ->
                     %% Reigster metrics
-                    Prefix = lists:flatten(io_lib:format("exhook.~s.", [Name])),
+                    Prefix = lists:flatten(
+                               io_lib:format("exhook.~s.", [Name])),
                     ensure_metrics(Prefix, HookSpecs),
+                    %% Ensure hooks
+                    ensure_hooks(HookSpecs),
                     {ok, #server{name = Name,
                                  options = Opts0,
                                  channel = _ChannPoolPid,
@@ -126,8 +133,9 @@ channel_opts(Opts) ->
     {SvrAddr, ClientOpts}.
 
 -spec unload(server()) -> ok.
-unload(#server{name = Name}) ->
+unload(#server{name = Name, hookspec = HookSpecs}) ->
     _ = do_deinit(Name),
+    _ = may_unload_hooks(HookSpecs),
     _ = emqx_exhook_sup:stop_grpc_client_channel(Name),
     ok.
 
@@ -177,6 +185,31 @@ ensure_metrics(Prefix, HookSpecs) ->
             || Hookpoint <- maps:keys(HookSpecs)],
     lists:foreach(fun emqx_metrics:ensure/1, Keys).
 
+ensure_hooks(HookSpecs) ->
+    lists:foreach(fun(Hookpoint) ->
+        case ets:lookup(?REGISTRAY, Hookpoint) of
+            [] ->
+                ?LOG(warning, "Hoook ~s not found in registray", [Hookpoint]);
+            [{Hookpoint, {M, F, A}, _}] ->
+                emqx_hooks:put(Hookpoint, {M, F, A}),
+                ets:update_counter(?REGISTRAY, Hookpoint, {3, 1})
+        end
+    end, maps:keys(HookSpecs)).
+
+may_unload_hooks(HookSpecs) ->
+    lists:foreach(fun(Hookpoint) ->
+        case ets:update_counter(?REGISTRAY, Hookpoint, {3, -1}) of
+            Cnt when Cnt =< 0 ->
+                case ets:lookup(?REGISTRAY, Hookpoint) of
+                    [{Hookpoint, {M, F, _A}, _}] ->
+                        emqx_hooks:del(Hookpoint, {M, F});
+                    _ -> ok
+                end,
+                ets:delete(?REGISTRAY, Hookpoint);
+            _ -> ok
+        end
+    end, maps:keys(HookSpecs)).
+
 format(#server{name = Name, hookspec = Hooks}) ->
     io_lib:format("name=~p, hooks=~0p", [Name, Hooks]).