Correctly handle multiple load attempts when on_load is pending

If an on_load function had not yet returned, the code server
would only correctly handle calls to code:ensure_loaded/1. That
is, each process that called code:ensured_loaded/1 for the
module in question would be suspended until the on_load function
had returned.

The code server handled calls to code:load_binary/1, code:load_file/1,
and code:load_abs/1 in the same way as for code:ensure_loaded/1.
That means that if call to one those functions attempted to load
*different* code for the module, that code would not get loaded.

Note that code:finish_loading/1 (code:atomic_load/1) will still return
{error,pending_on_load} if there is a pending on_load function for one
of the modules that are about to be loaded. The reason is that
code:finish_loading/1 is meant to either succeed directly, or fail
quickly if there is any problem.

2 files changed, 133 insertions, 56 deletions

diff --git a/lib/kernel/src/code_server.erl b/lib/kernel/src/code_server.erl
index 10264382eb..6174136507 100644
--- a/lib/kernel/src/code_server.erl
+++ b/lib/kernel/src/code_server.erl
@@ -32,8 +32,12 @@
 
 -import(lists, [foreach/2]).
 
+-type on_load_action() ::
+	fun((term(), state()) -> {'reply',term(),state()} |
+				 {'noreply',state()}).
+
 -type on_load_item() :: {{pid(),reference()},module(),
-			 file:name_all(),[pid()]}.
+			 [{pid(),on_load_action()}]}.
 
 -record(state, {supervisor :: pid(),
 		root :: file:name_all(),
@@ -304,7 +308,7 @@ handle_call({ensure_loaded,Mod}, From, St) when is_atom(Mod) ->
 	true ->
 	    {reply,{module,Mod},St};
 	false when St#state.mode =:= interactive ->
-	    load_file(Mod, From, St);
+	    ensure_loaded(Mod, From, St);
 	false ->
 	    {reply,{error,embedded},St}
     end;
@@ -1084,13 +1088,11 @@ load_abs(File, Mod, From, St) ->
 	    {reply,{error,nofile},St}
     end.
 
-try_load_module(File, Mod, Bin, From, St0) ->
-    case pending_on_load(Mod, From, St0) of
-	no ->
-	    try_load_module_1(File, Mod, Bin, From, St0);
-	{yes,St} ->
-	    {noreply,St}
-    end.
+try_load_module(File, Mod, Bin, From, St) ->
+    Action = fun(_, S) ->
+		     try_load_module_1(File, Mod, Bin, From, S)
+	     end,
+    handle_pending_on_load(Action, Mod, From, St).
 
 try_load_module_1(File, Mod, Bin, From, #state{moddb=Db}=St) ->
     case is_sticky(Mod, Db) of
@@ -1117,19 +1119,19 @@ try_load_module_2(File, Mod, Bin, From, Architecture,
             {reply,ok,St}
     end.
 
-try_load_module_3(File, Mod, Bin, From, Architecture,
-                  #state{moddb=Db}=St) ->
-    case erlang:load_module(Mod, Bin) of
-        {module,Mod} = Module ->
-            ets:insert(Db, {Mod,File}),
-            post_beam_load([Mod], Architecture, St),
-            {reply,Module,St};
-        {error,on_load} ->
-            handle_on_load(Mod, File, From, St);
-        {error,What} = Error ->
-            error_msg("Loading of ~ts failed: ~p\n", [File, What]),
-            {reply,Error,St}
-    end.
+try_load_module_3(File, Mod, Bin, From, Architecture, St0) ->
+    Action = fun({module,_}=Module, #state{moddb=Db}=S) ->
+		     ets:insert(Db, {Mod,File}),
+		     post_beam_load([Mod], Architecture, S),
+		     {reply,Module,S};
+		({error,on_load_failure}=Error, S) ->
+		     {reply,Error,S};
+		({error,What}=Error, S) ->
+		     error_msg("Loading of ~ts failed: ~p\n", [File, What]),
+		     {reply,Error,S}
+	     end,
+    Res = erlang:load_module(Mod, Bin),
+    handle_on_load(Res, Action, Mod, From, St0).
 
 hipe_result_to_status(Result, #state{moddb=Db}) ->
     case Result of
@@ -1154,11 +1156,22 @@ int_list([H|T]) when is_integer(H) -> int_list(T);
 int_list([_|_])                    -> false;
 int_list([])                       -> true.
 
+ensure_loaded(Mod, From, St0) ->
+    Action = fun(_, S) ->
+		     case erlang:module_loaded(Mod) of
+			 true ->
+			     {reply,{module,Mod},S};
+			 false ->
+			     load_file_1(Mod, From, S)
+		     end
+	     end,
+    handle_pending_on_load(Action, Mod, From, St0).
+
 load_file(Mod, From, St0) ->
-    case pending_on_load(Mod, From, St0) of
-	no -> load_file_1(Mod, From, St0);
-	{yes,St} -> {noreply,St}
-    end.
+    Action = fun(_, S) ->
+		     load_file_1(Mod, From, S)
+	     end,
+    handle_pending_on_load(Action, Mod, From, St0).
 
 load_file_1(Mod, From, #state{path=Path}=St) ->
     case mod_to_bin(Path, Mod) of
@@ -1308,55 +1321,56 @@ run([F|Fs], Data0) ->
 %% The on_load functionality.
 %% -------------------------------------------------------
 
-handle_on_load(Mod, File, From, #state{on_load=OnLoad0}=St0) ->
+handle_on_load({error,on_load}, Action, Mod, From, St0) ->
+    #state{on_load=OnLoad0} = St0,
     Fun = fun() ->
 		  Res = erlang:call_on_load_function(Mod),
 		  exit(Res)
 	  end,
     PidRef = spawn_monitor(Fun),
-    OnLoad = [{PidRef,Mod,File,[From]}|OnLoad0],
+    PidAction = {From,Action},
+    OnLoad = [{PidRef,Mod,[PidAction]}|OnLoad0],
     St = St0#state{on_load=OnLoad},
-    {noreply,St}.
+    {noreply,St};
+handle_on_load(Res, Action, _, _, St) ->
+    Action(Res, St).
 
-pending_on_load(_, _, #state{on_load=[]}) ->
-    no;
-pending_on_load(Mod, From, #state{on_load=OnLoad0}=St) ->
+handle_pending_on_load(Action, Mod, From, #state{on_load=OnLoad0}=St) ->
     case lists:keyfind(Mod, 2, OnLoad0) of
 	false ->
-	    no;
-	{{From,_Ref},Mod,_File,_Pids} ->
+	    Action(ok, St);
+	{{From,_Ref},Mod,_Pids} ->
 	    %% The on_load function tried to make an external
 	    %% call to its own module. That would be a deadlock.
 	    %% Fail the call. (The call is probably from error_handler,
 	    %% and it will ignore the actual error reason and cause
 	    %% an undef execption.)
-	    _ = reply(From, {error,deadlock}),
-	    {yes,St};
-	{_,_,_,_} ->
-	    OnLoad = pending_on_load_1(Mod, From, OnLoad0),
-	    {yes,St#state{on_load=OnLoad}}
+	    {reply,{error,deadlock},St};
+	{_,_,_} ->
+	    OnLoad = handle_pending_on_load_1(Mod, {From,Action}, OnLoad0),
+	    {noreply,St#state{on_load=OnLoad}}
     end.
 
-pending_on_load_1(Mod, From, [{PidRef,Mod,File,Pids}|T]) ->
-    [{PidRef,Mod,File,[From|Pids]}|T];
-pending_on_load_1(Mod, From, [H|T]) ->
-    [H|pending_on_load_1(Mod, From, T)];
-pending_on_load_1(_, _, []) -> [].
+handle_pending_on_load_1(Mod, From, [{PidRef,Mod,Pids}|T]) ->
+    [{PidRef,Mod,[From|Pids]}|T];
+handle_pending_on_load_1(Mod, From, [H|T]) ->
+    [H|handle_pending_on_load_1(Mod, From, T)];
+handle_pending_on_load_1(_, _, []) -> [].
 
-finish_on_load(PidRef, OnLoadRes, #state{on_load=OnLoad0,moddb=Db}=State) ->
+finish_on_load(PidRef, OnLoadRes, #state{on_load=OnLoad0}=St0) ->
     case lists:keyfind(PidRef, 1, OnLoad0) of
 	false ->
 	    %% Since this process in general silently ignores messages
 	    %% it doesn't understand, it should also ignore a 'DOWN'
 	    %% message with an unknown reference.
-	    State;
-	{PidRef,Mod,File,WaitingPids} ->
-	    finish_on_load_1(Mod, File, OnLoadRes, WaitingPids, Db),
-	    OnLoad = [E || {R,_,_,_}=E <- OnLoad0, R =/= PidRef],
-	    State#state{on_load=OnLoad}
+	    St0;
+	{PidRef,Mod,Waiting} ->
+	    St = finish_on_load_1(Mod, OnLoadRes, Waiting, St0),
+	    OnLoad = [E || {R,_,_}=E <- OnLoad0, R =/= PidRef],
+	    St#state{on_load=OnLoad}
     end.
 
-finish_on_load_1(Mod, File, OnLoadRes, WaitingPids, Db) ->
+finish_on_load_1(Mod, OnLoadRes, Waiting, St) ->
     Keep = OnLoadRes =:= ok,
     erlang:finish_after_on_load(Mod, Keep),
     Res = case Keep of
@@ -1365,11 +1379,20 @@ finish_on_load_1(Mod, File, OnLoadRes, WaitingPids, Db) ->
 		  _ = erts_code_purger:purge(Mod),
 		  {error,on_load_failure};
 	      true ->
-		  ets:insert(Db, {Mod,File}),
 		  {module,Mod}
 	  end,
-    _ = [reply(Pid, Res) || Pid <- WaitingPids],
-    ok.
+    finish_on_load_2(Waiting, Res, St).
+
+finish_on_load_2([{Pid,Action}|T], Res, St0) ->
+    case Action(Res, St0) of
+	{reply,Rep,St} ->
+	    _ = reply(Pid, Rep),
+	    finish_on_load_2(T, Res, St);
+	{noreply,St} ->
+	    finish_on_load_2(T, Res, St)
+    end;
+finish_on_load_2([], _, St) ->
+    St.
 
 finish_on_load_report(_Mod, Atom) when is_atom(Atom) ->
     %% No error reports for atoms.
diff --git a/lib/kernel/test/code_SUITE.erl b/lib/kernel/test/code_SUITE.erl
index e238eea2de..8da67c89f8 100644
--- a/lib/kernel/test/code_SUITE.erl
+++ b/lib/kernel/test/code_SUITE.erl
@@ -35,7 +35,7 @@
 	 purge_stacktrace/1, mult_lib_roots/1, bad_erl_libs/1,
 	 code_archive/1, code_archive2/1, on_load/1, on_load_binary/1,
 	 on_load_embedded/1, on_load_errors/1, on_load_update/1,
-	 on_load_purge/1, on_load_self_call/1,
+	 on_load_purge/1, on_load_self_call/1, on_load_pending/1,
 	 big_boot_embedded/1,
 	 native_early_modules/1, get_mode/1,
 	 normalized_paths/1]).
@@ -65,7 +65,7 @@ all() ->
      purge_stacktrace, mult_lib_roots,
      bad_erl_libs, code_archive, code_archive2, on_load,
      on_load_binary, on_load_embedded, on_load_errors, on_load_update,
-     on_load_purge, on_load_self_call,
+     on_load_purge, on_load_self_call, on_load_pending,
      big_boot_embedded, native_early_modules, get_mode, normalized_paths].
 
 groups() ->
@@ -830,6 +830,12 @@ check_funs({'$M_EXPR',warning_msg,2},
 	   [{code_server,finish_on_load_report,2} | _]) -> 0;
 check_funs({'$M_EXPR','$F_EXPR',1},
 	   [{code_server,run,2}|_]) -> 0;
+check_funs({'$M_EXPR','$F_EXPR',2},
+	   [{code_server,handle_on_load,5}|_]) -> 0;
+check_funs({'$M_EXPR','$F_EXPR',2},
+	   [{code_server,handle_pending_on_load,4}|_]) -> 0;
+check_funs({'$M_EXPR','$F_EXPR',2},
+	   [{code_server,finish_on_load_2,3}|_]) -> 0;
 %% This is cheating! /raimo
 %%
 %% check_funs(This = {M,_,_}, Path) ->
@@ -1549,6 +1555,54 @@ on_load_do_load(Mod, Code) ->
 	Any -> Any
     end.
 
+on_load_pending(_Config) ->
+    Mod = ?FUNCTION_NAME,
+    Tree1 = ?Q(["-module('@Mod@').\n",
+		"-on_load(f/0).\n",
+		"f() ->\n",
+		"  register('@Mod@', self()),\n",
+		"  receive _ -> ok end.\n"]),
+    merl:print(Tree1),
+    {ok,Mod,Code1} = merl:compile(Tree1),
+
+    Tree2 = ?Q(["-module('@Mod@').\n",
+		"-export([t/0]).\n",
+		"t() -> ok.\n"]),
+    merl:print(Tree2),
+    {ok,Mod,Code2} = merl:compile(Tree2),
+
+    Self = self(),
+    {_,Ref1} =
+	spawn_monitor(fun() ->
+			      Self ! started1,
+			      {module,Mod} = code:load_binary(Mod, "", Code1)
+		      end),
+    receive started1 -> ok end,
+    timer:sleep(10),
+    {_,Ref2} =
+	spawn_monitor(fun() ->
+			      Self ! started2,
+			      {module,Mod} = code:load_binary(Mod, "", Code2),
+			      ok = Mod:t()
+		      end),
+    receive started2 -> ok end,
+    receive
+	Unexpected ->
+	    ct:fail({unexpected,Unexpected})
+    after 100 ->
+	    ok
+    end,
+    Mod ! go,
+    receive
+	{'DOWN',Ref1,process,_,normal} -> ok
+    end,
+    receive
+	{'DOWN',Ref2,process,_,normal} -> ok
+    end,
+    ok = Mod:t(),
+    ok.
+
+
 %% Test that the native code of early loaded modules is loaded.
 native_early_modules(Config) when is_list(Config) ->
     case erlang:system_info(hipe_architecture) of