Add code_server cache to prevent code:get_object_code abuse

In the case of high connection churn, code:get_object_code is called multiple times to find the Module:additional_authn_params method for the direct connection type. In the case of a missing module, the code_server process can be overrun with messages while checking the code path. This change caches the result so that future calls to a missing or bad module are not as time consuming

Fixes VESC-888

3 files changed, 99 insertions, 15 deletions

diff --git a/src/code_server_cache.erl b/src/code_server_cache.erl
new file mode 100644
index 0000000000..a7acdd7fdb
--- /dev/null
+++ b/src/code_server_cache.erl
@@ -0,0 +1,90 @@
+%% -*- erlang-indent-level: 4;indent-tabs-mode: nil -*-
+%% ex: ts=4 sw=4 et
+%% The contents of this file are subject to the Mozilla Public License
+%% Version 1.1 (the "License"); you may not use this file except in
+%% compliance with the License. You may obtain a copy of the License
+%% at http://www.mozilla.org/MPL/
+%%
+%% Software distributed under the License is distributed on an "AS IS"
+%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
+%% the License for the specific language governing rights and
+%% limitations under the License.
+%%
+%% The Original Code is RabbitMQ.
+%%
+%% The Initial Developer of the Original Code is GoPivotal, Inc.
+%% Copyright (c) 2007-2017 Pivotal Software, Inc.  All rights reserved.
+%%
+
+-module(code_server_cache).
+
+-behaviour(gen_server).
+
+%% API
+-export([start_link/0,
+         maybe_call_mfa/4]).
+
+%% gen_server callbacks
+-export([init/1,
+         handle_call/3,
+         handle_cast/2,
+         handle_info/2,
+         terminate/2,
+         code_change/3]).
+
+-record(state, {
+    modules = #{} :: #{atom() => boolean()}
+}).
+
+%% API
+start_link() ->
+    gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
+
+maybe_call_mfa(Module, Function, Args, Default) ->
+    gen_server:call(?MODULE, {maybe_call_mfa, {Module, Function, Args, Default}}).
+
+%% gen_server callbacks
+
+init([]) ->
+    {ok, #state{}}.
+
+handle_call({maybe_call_mfa, {Mod, _F, _A, _D} = MFA}, _From, #state{modules = ModuleMap} = State0) ->
+    Value = maps:get(Mod, ModuleMap, true),
+    {ok, Reply, State1} = handle_maybe_call_mfa(Value, MFA, State0),
+    {reply, Reply, State1};
+handle_call(_Request, _From, State) ->
+    {reply, ignored, State}.
+
+handle_cast(_Msg, State) ->
+    {noreply, State}.
+
+handle_info(_Info, State) ->
+    {noreply, State}.
+
+terminate(_Reason, _State) ->
+    ok.
+
+code_change(_OldVsn, State, _Extra) ->
+    {ok, State}.
+
+%% Internal functions
+
+handle_maybe_call_mfa(false, {_M, _F, _A, Default}, State) ->
+    {ok, Default, State};
+handle_maybe_call_mfa(true, {Module, Function, Args, Default}, State) ->
+    try
+        Reply = erlang:apply(Module, Function, Args),
+        {ok, Reply, State}
+    catch
+        error:undef ->
+            handle_maybe_call_mfa_error(Module, Default, State);
+        Err:Reason ->
+            rabbit_log:error("Calling ~p:~p failed: ~p:~p~n",
+                             [Module, Function, Err, Reason]),
+            handle_maybe_call_mfa_error(Module, Default, State)
+    end.
+
+handle_maybe_call_mfa_error(Module, Default, #state{modules = ModuleMap0} = State0) ->
+    ModuleMap1 = maps:put(Module, false, ModuleMap0),
+    State1 = State0#state{modules = ModuleMap1},
+    {ok, Default, State1}.
diff --git a/src/rabbit.erl b/src/rabbit.erl
index 394ae699e3..57aa825a1b 100644
--- a/src/rabbit.erl
+++ b/src/rabbit.erl
@@ -74,11 +74,17 @@
                     {requires,    database},
                     {enables,     external_infrastructure}]}).
 
+-rabbit_boot_step({code_server_cache,
+                   [{description, "code_server cache server"},
+                    {mfa,         {rabbit_sup, start_child, [code_server_cache]}},
+                    {requires,    rabbit_alarm},
+                    {enables,     file_handle_cache}]}).
+
 -rabbit_boot_step({file_handle_cache,
                    [{description, "file handle cache server"},
                     {mfa,         {rabbit, start_fhc, []}},
                     %% FHC needs memory monitor to be running
-                    {requires,    rabbit_alarm},
+                    {requires,    code_server_cache},
                     {enables,     worker_pool}]}).
 
 -rabbit_boot_step({worker_pool,
diff --git a/src/rabbit_direct.erl b/src/rabbit_direct.erl
index 074eed1f9a..c3d8e70efd 100644
--- a/src/rabbit_direct.erl
+++ b/src/rabbit_direct.erl
@@ -135,20 +135,8 @@ maybe_call_connection_info_module(Protocol, Creds, VHost, Pid, Infos) ->
     Module = rabbit_data_coercion:to_atom(string:to_lower(
         "rabbit_" ++ rabbit_data_coercion:to_list(Protocol) ++ "_connection_info")
     ),
-    case code:get_object_code(Module) of
-        {_Module, _Binary, _Filename} ->
-            try
-                Module:additional_authn_params(Creds, VHost, Pid, Infos)
-            catch
-                _:Reason ->
-                    rabbit_log:error("Calling ~p:additional_authn_params/4 failed:~p~n", [Module, Reason]),
-                    []
-            end;
-        error ->
-            [];
-        _ ->
-            []
-    end.
+    Args = [Creds, VHost, Pid, Infos],
+    code_server_cache:maybe_call_mfa(Module, additional_authn_params, Args, []).
 
 is_vhost_alive(VHost, {Username, _Password}, Pid) ->
     PrintedUsername = case Username of