Don't leak old code when loading a modules with an on_load function

Normally, calling code:delete/1 before re-loading the code for a
module is unnecessary but causes no problem.

But there will be be problems if the new code has an on_load function.
Code with an on_load function will always be loaded as old code
to allowed it to be easily purged if the on_load function would fail.
If the on_load function succeeds, the old and current code will be
swapped.

So in the scenario where code:delete/1 has been called explicitly,
there is old code but no current code. Loading code with an
on_load function will cause the reference to the old code to be
overwritten. That will at best cause a memory leak, and at worst
an emulator crash (especially if NIFs are involved).

To avoid that situation, we will put the code with the on_load
function in a special, third slot in Module.

ERL-240

2 files changed, 37 insertions, 2 deletions

diff --git a/erts/preloaded/src/erts_code_purger.erl b/erts/preloaded/src/erts_code_purger.erl
index ee4fcedd2d..28d71fd07e 100644
--- a/erts/preloaded/src/erts_code_purger.erl
+++ b/erts/preloaded/src/erts_code_purger.erl
@@ -22,7 +22,8 @@
 %% Purpose : Implement system process erts_code_purger
 %%           to handle code module purging.
 
--export([start/0, purge/1, soft_purge/1, pending_purge_lambda/3]).
+-export([start/0, purge/1, soft_purge/1, pending_purge_lambda/3,
+	 finish_after_on_load/2]).
 
 -spec start() -> term().
 start() ->
@@ -40,6 +41,11 @@ loop() ->
 	    Res = do_soft_purge(Mod),
 	    From ! {reply, soft_purge, Res, Ref};
 
+	{finish_after_on_load,{Mod,Keep},From,Ref}
+	      when is_atom(Mod), is_pid(From) ->
+	    Res = do_finish_after_on_load(Mod, Keep),
+	    From ! {reply, finish_after_on_load, Res, Ref};
+
 	{test_purge, Mod, From, Type, Ref} when is_atom(Mod), is_pid(From) ->
 	     do_test_purge(Mod, From, Type, Ref);
 
@@ -129,6 +135,35 @@ do_soft_purge(Mod) ->
 				       end)
     end.
 
+%% finish_after_on_load(Module, Keep)
+%% Finish after running on_load function. If Keep is false,
+%% purge the code for the on_load function.
+
+finish_after_on_load(Mod, Keep) ->
+    Ref = make_ref(),
+    erts_code_purger ! {finish_after_on_load, {Mod,Keep}, self(), Ref},
+    receive
+	{reply, finish_after_on_load, Result, Ref} ->
+	    Result
+    end.
+
+do_finish_after_on_load(Mod, Keep) ->
+    erlang:finish_after_on_load(Mod, Keep),
+    case Keep of
+	true ->
+	    ok;
+	false ->
+	    case erts_internal:purge_module(Mod, prepare_on_load) of
+		false ->
+		    true;
+		true ->
+		    _ = check_proc_code(erlang:processes(), Mod, true),
+		    true = erts_internal:purge_module(Mod, complete)
+	    end
+    end.
+
+
+
 %%
 %% check_proc_code(Pids, Mod, Hard) - Send asynchronous
 %%   requests to all processes to perform a check_process_code
diff --git a/erts/preloaded/src/erts_internal.erl b/erts/preloaded/src/erts_internal.erl
index 6229754c8c..6aae5ba38c 100644
--- a/erts/preloaded/src/erts_internal.erl
+++ b/erts/preloaded/src/erts_internal.erl
@@ -304,7 +304,7 @@ release_literal_area_switch() ->
 
 -spec purge_module(Module, Op) -> boolean() when
       Module :: module(),
-      Op :: 'prepare' | 'abort' | 'complete'.
+      Op :: 'prepare' | 'prepare_on_load' | 'abort' | 'complete'.
 purge_module(_Module, _Op) ->
     erlang:nif_error(undefined).