Merge branch 'bjorn/erts/on_load/ERL-240/OTP-13893' into maint

* bjorn/erts/on_load/ERL-240/OTP-13893:
  erts: Add nif_SUITE:t_on_load
  erts: Improve nif_SUITE:upgrade test
  Don't leak old code when loading a modules with an on_load function

14 files changed, 406 insertions, 30 deletions

diff --git a/erts/emulator/beam/atom.names b/erts/emulator/beam/atom.names
index 9dae67cb2d..2b66bf6f4e 100644
--- a/erts/emulator/beam/atom.names
+++ b/erts/emulator/beam/atom.names
@@ -505,6 +505,7 @@ atom port_limit
 atom port_op
 atom positive
 atom prepare
+atom prepare_on_load
 atom print
 atom priority
 atom private
diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index dddcfbb77d..5969197168 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -53,6 +53,7 @@ static struct {
     ErlFunEntry *def_funs[10];
     Uint fe_size;
     Uint fe_ix;
+    struct erl_module_instance saved_old;
 } purge_state;
 
 Process *erts_code_purger = NULL;
@@ -105,6 +106,8 @@ init_purge_state(void)
     purge_state.fe_size = sizeof(purge_state.def_funs);
     purge_state.fe_size /= sizeof(purge_state.def_funs[0]);
     purge_state.fe_ix = 0;
+
+    purge_state.saved_old.code_hdr = 0;
 }
 
 void
@@ -739,10 +742,13 @@ BIF_RETTYPE call_on_load_function_1(BIF_ALIST_1)
 {
     Module* modp = erts_get_module(BIF_ARG_1, erts_active_code_ix());
 
-    if (modp && modp->old.code_hdr) {
-	BIF_TRAP_CODE_PTR_0(BIF_P, modp->old.code_hdr->on_load_function_ptr);
+    if (!modp || !modp->on_load) {
+	BIF_ERROR(BIF_P, BADARG);
     }
-    else {
+    if (modp->on_load->code_hdr) {
+	BIF_TRAP_CODE_PTR_0(BIF_P,
+			    modp->on_load->code_hdr->on_load_function_ptr);
+    } else {
 	BIF_ERROR(BIF_P, BADARG);
     }
 }
@@ -765,14 +771,14 @@ BIF_RETTYPE finish_after_on_load_2(BIF_ALIST_2)
     code_ix = erts_active_code_ix();
     modp = erts_get_module(BIF_ARG_1, code_ix);
 
-    if (!modp || !modp->old.code_hdr) {
+    if (!modp || !modp->on_load || !modp->on_load->code_hdr) {
     error:
 	erts_smp_thr_progress_unblock();
         erts_smp_proc_lock(BIF_P, ERTS_PROC_LOCK_MAIN);
 	erts_release_code_write_permission();
 	BIF_ERROR(BIF_P, BADARG);
     }
-    if (modp->old.code_hdr->on_load_function_ptr == NULL) {
+    if (modp->on_load->code_hdr->on_load_function_ptr == NULL) {
 	goto error;
     }
     if (BIF_ARG_2 != am_false && BIF_ARG_2 != am_true) {
@@ -781,14 +787,17 @@ BIF_RETTYPE finish_after_on_load_2(BIF_ALIST_2)
 
     if (BIF_ARG_2 == am_true) {
 	int i;
-	struct erl_module_instance t;
 
 	/*
-	 * Swap old and new code.
+	 * Make the code with the on_load function current.
 	 */
-	t = modp->curr;
-	modp->curr = modp->old;
-	modp->old = t;
+
+	if (modp->curr.code_hdr) {
+	    modp->old = modp->curr;
+	}
+	modp->curr = *modp->on_load;
+	erts_free(ERTS_ALC_T_PREPARED_CODE, modp->on_load);
+	modp->on_load = 0;
 
 	/*
 	 * The on_load function succeded. Fix up export entries.
@@ -1731,7 +1740,8 @@ BIF_RETTYPE erts_internal_purge_module_2(BIF_ALIST_2)
 
     switch (BIF_ARG_2) {
 
-    case am_prepare: {
+    case am_prepare:
+    case am_prepare_on_load: {
 	/*
 	 * Prepare for purge by marking all fun
 	 * entries referring to the code to purge
@@ -1756,7 +1766,22 @@ BIF_RETTYPE erts_internal_purge_module_2(BIF_ALIST_2)
 	    /*
 	     * Any code to purge?
 	     */
-	    erts_rlock_old_code(code_ix);
+
+	    if (BIF_ARG_2 == am_prepare_on_load) {
+                erts_rwlock_old_code(code_ix);
+	    } else {
+		erts_rlock_old_code(code_ix);
+	    }
+
+	    if (BIF_ARG_2 == am_prepare_on_load) {
+		ASSERT(modp->on_load);
+		ASSERT(modp->on_load->code_hdr);
+		purge_state.saved_old = modp->old;
+		modp->old = *modp->on_load;
+		erts_free(ERTS_ALC_T_PREPARED_CODE, (void *) modp->on_load);
+		modp->on_load = 0;
+	    }
+
 	    if (!modp->old.code_hdr)
 		res = am_false;
 	    else {
@@ -1774,7 +1799,12 @@ BIF_RETTYPE erts_internal_purge_module_2(BIF_ALIST_2)
 		ERTS_SET_COPY_LITERAL_AREA(modp->old.code_hdr->literal_area);
 #endif
 	    }
-	    erts_runlock_old_code(code_ix);
+
+            if (BIF_ARG_2 == am_prepare_on_load) {
+                erts_rwunlock_old_code(code_ix);
+	    } else {
+                erts_runlock_old_code(code_ix);
+	    }
 	}
 	
 #ifndef ERTS_SMP
@@ -1911,8 +1941,14 @@ BIF_RETTYPE erts_internal_purge_module_2(BIF_ALIST_2)
 		modp->old.code_length = 0;
 		modp->old.catches = BEAM_CATCHES_NIL;
 		erts_remove_from_ranges(code);
+
 		ERTS_BIF_PREP_RET(ret, am_true);
 	    }
+
+	    if (purge_state.saved_old.code_hdr) {
+		modp->old = purge_state.saved_old;
+		purge_state.saved_old.code_hdr = 0;
+	    }
 	    erts_rwunlock_old_code(code_ix);
 	}
 	if (is_blocking) {
diff --git a/erts/emulator/beam/beam_load.c b/erts/emulator/beam/beam_load.c
index f63addb309..0afdedf6c2 100644
--- a/erts/emulator/beam/beam_load.c
+++ b/erts/emulator/beam/beam_load.c
@@ -833,11 +833,20 @@ erts_finish_loading(Binary* magic, Process* c_p,
 
     size = stp->loaded_size;
     erts_total_code_size += size;
-    if (stp->on_load) {
-	inst_p = &mod_tab_p->old;
-    } else {
+
+    if (!stp->on_load) {
 	inst_p = &mod_tab_p->curr;
+    } else {
+	mod_tab_p->on_load =
+	    (struct erl_module_instance *)
+	    erts_alloc(ERTS_ALC_T_PREPARED_CODE,
+		       sizeof(struct erl_module_instance));
+	inst_p = mod_tab_p->on_load;
+	inst_p->nif = 0;
+        inst_p->num_breakpoints = 0;
+        inst_p->num_traced_exports = 0;
     }
+
     inst_p->code_hdr = stp->hdr;
     inst_p->code_length = size;
 
diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 559b4017e7..3a547982da 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -3203,18 +3203,20 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
     if (init_func != NULL)
       handle = init_func;
 
+    this_mi = &module_p->curr;
+    prev_mi = &module_p->old;
     if (in_area(caller, module_p->old.code_hdr, module_p->old.code_length)) {
-	if (module_p->old.code_hdr->on_load_function_ptr) {
-	    this_mi = &module_p->old;
+	ret = load_nif_error(BIF_P, "old_code", "Calling load_nif from old "
+			     "module '%T' not allowed", mod_atom);
+	goto error;
+    } else if (module_p->on_load) {
+	ASSERT(module_p->on_load->code_hdr->on_load_function_ptr);
+	if (module_p->curr.code_hdr) {
 	    prev_mi = &module_p->curr;
 	} else {
-	    ret = load_nif_error(BIF_P, "old_code", "Calling load_nif from old "
-				 "module '%T' not allowed", mod_atom);
-	    goto error;
+	    prev_mi = &module_p->old;
 	}
-    } else {
-	this_mi = &module_p->curr;
-	prev_mi = &module_p->old;
+	this_mi = module_p->on_load;
     }
 
     if (init_func == NULL &&
diff --git a/erts/emulator/beam/module.c b/erts/emulator/beam/module.c
index 3eb11f1693..4f36377450 100644
--- a/erts/emulator/beam/module.c
+++ b/erts/emulator/beam/module.c
@@ -85,6 +85,7 @@ static Module* module_alloc(Module* tmpl)
     obj->old.num_breakpoints  = 0;
     obj->curr.num_traced_exports = 0;
     obj->old.num_traced_exports = 0;
+    obj->on_load = 0;
     return obj;
 }
 
@@ -201,6 +202,7 @@ void module_start_staging(void)
 
 	dst_mod->curr = src_mod->curr;
 	dst_mod->old = src_mod->old;
+	dst_mod->on_load = src_mod->on_load;
     }
 
     /*
@@ -214,6 +216,7 @@ void module_start_staging(void)
 
 	dst_mod->curr = src_mod->curr;
 	dst_mod->old = src_mod->old;
+	dst_mod->on_load = src_mod->on_load;
     }
     newsz = index_table_sz(dst);
     erts_smp_atomic_add_nob(&tot_module_bytes, (newsz - oldsz));
diff --git a/erts/emulator/beam/module.h b/erts/emulator/beam/module.h
index 5a60bc90d9..1c1afc8461 100644
--- a/erts/emulator/beam/module.h
+++ b/erts/emulator/beam/module.h
@@ -39,6 +39,7 @@ typedef struct erl_module {
 
     struct erl_module_instance curr;
     struct erl_module_instance old; /* protected by "old_code" rwlock */
+    struct erl_module_instance* on_load;
 } Module; 
 
 Module* erts_get_module(Eterm mod, ErtsCodeIndex code_ix);
diff --git a/erts/emulator/test/nif_SUITE.erl b/erts/emulator/test/nif_SUITE.erl
index 8df2733fac..9c1694fa8a 100644
--- a/erts/emulator/test/nif_SUITE.erl
+++ b/erts/emulator/test/nif_SUITE.erl
@@ -29,6 +29,7 @@
 -export([all/0, suite/0,
 	 init_per_testcase/2, end_per_testcase/2,
          basic/1, reload/1, upgrade/1, heap_frag/1,
+         t_on_load/1,
 	 types/1, many_args/1, binaries/1, get_string/1, get_atom/1,
 	 maps/1,
 	 api_macros/1,
@@ -68,6 +69,7 @@ suite() -> [{ct_hooks,[ts_install_cth]}].
 
 all() -> 
     [basic, reload, upgrade, heap_frag, types, many_args,
+     t_on_load,
      binaries, get_string, get_atom, maps, api_macros, from_array,
      iolist_as_binary, resource, resource_binary,
      resource_takeover, threading, send, send2, send3,
@@ -83,10 +85,19 @@ all() ->
      nif_port_command,
      nif_snprintf].
 
+init_per_testcase(t_on_load, Config) ->
+    ets:new(nif_SUITE, [named_table]),
+    Config;
 init_per_testcase(_Case, Config) ->
     Config.
 
+end_per_testcase(t_on_load, _Config) ->
+    ets:delete(nif_SUITE),
+    testcase_cleanup();
 end_per_testcase(_Func, _Config) ->
+    testcase_cleanup().
+
+testcase_cleanup() ->
     P1 = code:purge(nif_mod),
     Del = code:delete(nif_mod),
     P2 = code:purge(nif_mod),
@@ -174,11 +185,33 @@ upgrade(Config) when is_list(Config) ->
     true = erlang:delete_module(nif_mod),
     [] = nif_mod_call_history(),    
 
+    %% Repeat upgrade again but from old (deleted) instance
+    {module,nif_mod} = erlang:load_module(nif_mod,Bin),
+    undefined = nif_mod:lib_version(),
+    1 = call(Pid,lib_version),
+    [{lib_version,1,9,109}] = nif_mod_call_history(),
+
+    ok = nif_mod:load_nif_lib(Config, 1),
+    1 = nif_mod:lib_version(),
+    [{upgrade,1,10,110},{lib_version,1,11,111}] = nif_mod_call_history(),
+
+    upgraded = call(Pid,upgrade),
+    false = check_process_code(Pid, nif_mod),
+    true = erlang:purge_module(nif_mod),
+    [{unload,1,12,112}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    [{lib_version,1,13,113}] = nif_mod_call_history(),
+
+    true = erlang:delete_module(nif_mod),
+    [] = nif_mod_call_history(),
+
+
     Pid ! die,
     {'DOWN', MRef, process, Pid, normal} = receive_any(),
     false = check_process_code(Pid, nif_mod),
     true = erlang:purge_module(nif_mod),
-    [{unload,1,9,109}] = nif_mod_call_history(),    
+    [{unload,1,14,114}] = nif_mod_call_history(),
 
     %% Module upgrade with different lib version
     {module,nif_mod} = erlang:load_module(nif_mod,Bin),
@@ -215,17 +248,170 @@ upgrade(Config) when is_list(Config) ->
     true = erlang:delete_module(nif_mod),
     [] = nif_mod_call_history(),    
 
+
+    %% Reverse upgrade but from old (deleted) instance
+    {module,nif_mod} = erlang:load_module(nif_mod,Bin),
+    undefined = nif_mod:lib_version(),
+    [] = nif_mod_call_history(),
+    2 = call(Pid2,lib_version),
+    [{lib_version,2,4,204}] = nif_mod_call_history(),
+
+    ok = nif_mod:load_nif_lib(Config, 1),
+    1 = nif_mod:lib_version(),
+    [{upgrade,1,1,101},{lib_version,1,2,102}] = nif_mod_call_history(),
+
+    2 = call(Pid2,lib_version),
+    [{lib_version,2,5,205}] = nif_mod_call_history(),
+
+    upgraded = call(Pid2,upgrade),
+    false = check_process_code(Pid2, nif_mod),
+    true = erlang:purge_module(nif_mod),
+    [{unload,2,6,206}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    [{lib_version,1,3,103}] = nif_mod_call_history(),
+
+    true = erlang:delete_module(nif_mod),
+    [] = nif_mod_call_history(),
+
+
     Pid2 ! die,
     {'DOWN', MRef2, process, Pid2, normal} = receive_any(),
     false= check_process_code(Pid2, nif_mod),
     true = erlang:purge_module(nif_mod),
-    [{unload,2,4,204}] = nif_mod_call_history(),    
+    [{unload,1,4,104}] = nif_mod_call_history(),
 
     true = lists:member(?MODULE, erlang:system_info(taints)),
     true = lists:member(nif_mod, erlang:system_info(taints)),
     verify_tmpmem(TmpMem),
     ok.
 
+%% Test loading/upgrade in on_load
+t_on_load(Config) when is_list(Config) ->
+    TmpMem = tmpmem(),
+    ensure_lib_loaded(Config),
+
+    Data = proplists:get_value(data_dir, Config),
+    File = filename:join(Data, "nif_mod"),
+    {ok,nif_mod,Bin} = compile:file(File, [binary,return_errors,
+                                           {d,'USE_ON_LOAD'}]),
+
+    %% Use ETS to tell nif_mod:on_load what to do
+    ets:insert(nif_SUITE, {data_dir, Data}),
+    ets:insert(nif_SUITE, {lib_version, 1}),
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    hold_nif_mod_priv_data(nif_mod:get_priv_data_ptr()),
+    [{load,1,1,101},{get_priv_data_ptr,1,2,102}] = nif_mod_call_history(),
+
+    {Pid,MRef} = nif_mod:start(),
+    1 = call(Pid,lib_version),
+    [{lib_version,1,3,103}] = nif_mod_call_history(),
+
+    %% Module upgrade with same lib-version
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    1 = nif_mod:lib_version(),
+    1 = call(Pid,lib_version),
+    [{upgrade,1,4,104},{lib_version,1,5,105},{lib_version,1,6,106}] = nif_mod_call_history(),
+
+    upgraded = call(Pid,upgrade),
+    false = check_process_code(Pid, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,1,7,107}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    [{lib_version,1,8,108}] = nif_mod_call_history(),
+
+    true = code:delete(nif_mod),
+    [] = nif_mod_call_history(),
+
+    %% Repeat upgrade again but from old (deleted) instance
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    [{upgrade,1,9,109}] = nif_mod_call_history(),
+    1 = nif_mod:lib_version(),
+    1 = call(Pid,lib_version),
+    [{lib_version,1,10,110},{lib_version,1,11,111}] = nif_mod_call_history(),
+
+    upgraded = call(Pid,upgrade),
+    false = check_process_code(Pid, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,1,12,112}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    [{lib_version,1,13,113}] = nif_mod_call_history(),
+
+    true = code:delete(nif_mod),
+    [] = nif_mod_call_history(),
+
+
+    Pid ! die,
+    {'DOWN', MRef, process, Pid, normal} = receive_any(),
+    false = check_process_code(Pid, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,1,14,114}] = nif_mod_call_history(),
+
+    %% Module upgrade with different lib version
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    hold_nif_mod_priv_data(nif_mod:get_priv_data_ptr()),
+    [{load,1,1,101},{get_priv_data_ptr,1,2,102}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    {Pid2,MRef2} = nif_mod:start(),
+    1 = call(Pid2,lib_version),
+    [{lib_version,1,3,103},{lib_version,1,4,104}] = nif_mod_call_history(),
+
+    true = ets:insert(nif_SUITE,{lib_version,2}),
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    [{upgrade,2,1,201}] = nif_mod_call_history(),
+
+    2 = nif_mod:lib_version(),
+    1 = call(Pid2,lib_version),
+    [{lib_version,2,2,202},{lib_version,1,5,105}] = nif_mod_call_history(),
+
+    upgraded = call(Pid2,upgrade),
+    false = check_process_code(Pid2, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,1,6,106}] = nif_mod_call_history(),
+
+    2 = nif_mod:lib_version(),
+    2 = call(Pid2,lib_version),
+    [{lib_version,2,3,203},{lib_version,2,4,204}] = nif_mod_call_history(),
+
+    true = code:delete(nif_mod),
+    [] = nif_mod_call_history(),
+
+    %% Reverse upgrade but from old (deleted) instance
+    ets:insert(nif_SUITE,{lib_version,1}),
+    {module,nif_mod} = code:load_binary(nif_mod,File,Bin),
+    [{upgrade,1,1,101}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    2 = call(Pid2,lib_version),
+    [{lib_version,1,2,102},{lib_version,2,5,205}] = nif_mod_call_history(),
+
+    upgraded = call(Pid2,upgrade),
+    false = check_process_code(Pid2, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,2,6,206}] = nif_mod_call_history(),
+
+    1 = nif_mod:lib_version(),
+    [{lib_version,1,3,103}] = nif_mod_call_history(),
+
+    true = code:delete(nif_mod),
+    [] = nif_mod_call_history(),
+
+
+    Pid2 ! die,
+    {'DOWN', MRef2, process, Pid2, normal} = receive_any(),
+    false= check_process_code(Pid2, nif_mod),
+    true = code:soft_purge(nif_mod),
+    [{unload,1,4,104}] = nif_mod_call_history(),
+
+    true = lists:member(?MODULE, erlang:system_info(taints)),
+    true = lists:member(nif_mod, erlang:system_info(taints)),
+    verify_tmpmem(TmpMem),
+    ok.
+
+
 %% Test NIF building heap fragments
 heap_frag(Config) when is_list(Config) ->    
     TmpMem = tmpmem(),
diff --git a/erts/emulator/test/nif_SUITE_data/nif_mod.erl b/erts/emulator/test/nif_SUITE_data/nif_mod.erl
index eec1bb8858..1fcc33faa4 100644
--- a/erts/emulator/test/nif_SUITE_data/nif_mod.erl
+++ b/erts/emulator/test/nif_SUITE_data/nif_mod.erl
@@ -29,6 +29,16 @@
 
 -define(nif_stub,nif_stub_error(?LINE)).
 
+-ifdef(USE_ON_LOAD).
+-on_load(on_load/0).
+
+on_load() ->
+    [{data_dir, Path}] = ets:lookup(nif_SUITE, data_dir),
+    [{lib_version, Ver}] = ets:lookup(nif_SUITE, lib_version),
+    erlang:load_nif(filename:join(Path,libname(Ver)), []).
+
+-endif.
+
 load_nif_lib(Config, Ver) ->
     load_nif_lib(Config, Ver, []).
 
diff --git a/erts/preloaded/ebin/erts_code_purger.beam b/erts/preloaded/ebin/erts_code_purger.beam
index a1eb126098..6956eee740 100644
--- a/erts/preloaded/ebin/erts_code_purger.beam
+++ b/erts/preloaded/ebin/erts_code_purger.beam
diff --git a/erts/preloaded/ebin/erts_internal.beam b/erts/preloaded/ebin/erts_internal.beam
index 22817be8f4..227b62b7d3 100644
--- a/erts/preloaded/ebin/erts_internal.beam
+++ b/erts/preloaded/ebin/erts_internal.beam
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).
 
diff --git a/lib/kernel/src/code_server.erl b/lib/kernel/src/code_server.erl
index 29307bc696..59b26176bf 100644
--- a/lib/kernel/src/code_server.erl
+++ b/lib/kernel/src/code_server.erl
@@ -1392,11 +1392,10 @@ finish_on_load(PidRef, OnLoadRes, #state{on_load=OnLoad0}=St0) ->
 
 finish_on_load_1(Mod, OnLoadRes, Waiting, St) ->
     Keep = OnLoadRes =:= ok,
-    erlang:finish_after_on_load(Mod, Keep),
+    erts_code_purger:finish_after_on_load(Mod, Keep),
     Res = case Keep of
 	      false ->
 		  _ = finish_on_load_report(Mod, OnLoadRes),
-		  _ = erts_code_purger:purge(Mod),
 		  {error,on_load_failure};
 	      true ->
 		  {module,Mod}
diff --git a/lib/kernel/test/code_SUITE.erl b/lib/kernel/test/code_SUITE.erl
index 8da67c89f8..c5167efa56 100644
--- a/lib/kernel/test/code_SUITE.erl
+++ b/lib/kernel/test/code_SUITE.erl
@@ -36,6 +36,7 @@
 	 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_pending/1,
+	 on_load_deleted/1,
 	 big_boot_embedded/1,
 	 native_early_modules/1, get_mode/1,
 	 normalized_paths/1]).
@@ -66,6 +67,7 @@ all() ->
      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_pending,
+     on_load_deleted,
      big_boot_embedded, native_early_modules, get_mode, normalized_paths].
 
 groups() ->
@@ -1602,6 +1604,98 @@ on_load_pending(_Config) ->
     ok = Mod:t(),
     ok.
 
+on_load_deleted(_Config) ->
+    Mod = ?FUNCTION_NAME,
+
+    R0 = fun() ->
+		 Tree = ?Q(["-module('@Mod@').\n",
+			    "-on_load(f/0).\n",
+			    "f() -> ok.\n"]),
+		 merl:print(Tree),
+		 {ok,Mod,Code} = merl:compile(Tree),
+		 {module,Mod} = code:load_binary(Mod, "", Code)
+	 end,
+    delete_before_reload(Mod, R0),
+    delete_before_reload(Mod, R0),
+
+    R1 = fun() ->
+		 Tree = ?Q(["-module('@Mod@').\n",
+			    "-on_load(f/0).\n",
+			    "f() -> fail.\n"]),
+		 merl:print(Tree),
+		 {ok,Mod,Code} = merl:compile(Tree),
+		 {error,on_load_failure} = code:load_binary(Mod, "", Code)
+	 end,
+    delete_before_reload(Mod, R1),
+    delete_before_reload(Mod, R1),
+
+    OtherMod = list_to_atom(lists:concat([Mod,"_42"])),
+    OtherTree = ?Q(["-module('@OtherMod@').\n"]),
+    merl:print(OtherTree),
+    {ok,OtherMod,OtherCode} = merl:compile(OtherTree),
+
+    R2 = fun() ->
+		 RegName = 'on_load__registered_name',
+		 Tree = ?Q(["-module('@Mod@').\n",
+			    "-on_load(f/0).\n",
+			    "f() ->\n",
+			    "  register('@RegName@', self()),\n",
+			    "  receive _ -> ok end.\n"]),
+		 merl:print(Tree),
+		 {ok,Mod,Code} = merl:compile(Tree),
+		 spawn(fun() ->
+			       {module,Mod} = code:load_binary(Mod, "", Code)
+		       end),
+		 receive after 1 -> ok end,
+		 {module,OtherMod} = code:load_binary(OtherMod, "",
+						      OtherCode),
+		 RegName ! stop
+	 end,
+    delete_before_reload(Mod, R2),
+
+    ok.
+
+delete_before_reload(Mod, Reload) ->
+    false = check_old_code(Mod),
+
+    Tree1 = ?Q(["-module('@Mod@').\n",
+		"-export([f/1]).\n",
+		"f(Parent) ->\n",
+		"  register('@Mod@', self()),\n",
+		"  Parent ! started,\n",
+		"  receive _ -> ok end.\n"]),
+    merl:print(Tree1),
+    {ok,Mod,Code1} = merl:compile(Tree1),
+
+    Self = self(),
+    spawn(fun() ->
+		  {module,Mod} = code:load_binary(Mod, "", Code1),
+		  Mod:f(Self)
+	  end),
+    receive started -> ok end,
+
+    true = code:delete(Mod),
+    true = check_old_code(Mod),
+
+    Reload(),
+
+    %% When loading the the module with the -on_load() function,
+    %% the reference to the old code would be lost. Make sure that
+    %% the old code is remembered and is still preventing the
+    %% purge.
+    false = code:soft_purge(Mod),
+
+    %% Get rid of the old code.
+    Mod ! stop,
+    receive after 1 -> ok end,
+    true = code:soft_purge(Mod),
+
+    %% Unload the version of the module with the -on_load() function.
+    true = code:delete(Mod),
+    true = code:soft_purge(Mod),
+
+    ok.
+
 
 %% Test that the native code of early loaded modules is loaded.
 native_early_modules(Config) when is_list(Config) ->