Reimplement -on_load()

Load the module as old code; swap old and new code if the
-on_load function succeeds. That way, a failed update attempt
for a module that has an -on_load function will preserve the
previous version of the code.

5 files changed, 212 insertions, 63 deletions

diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index 87508dcf5f..b6f5d66166 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -625,8 +625,8 @@ 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->curr.code_hdr) {
-	BIF_TRAP_CODE_PTR_0(BIF_P, modp->curr.code_hdr->on_load_function_ptr);
+    if (modp && modp->old.code_hdr) {
+	BIF_TRAP_CODE_PTR_0(BIF_P, modp->old.code_hdr->on_load_function_ptr);
     }
     else {
 	BIF_ERROR(BIF_P, BADARG);
@@ -651,14 +651,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->curr.code_hdr) {
+    if (!modp || !modp->old.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->curr.code_hdr->on_load_function_ptr == NULL) {
+    if (modp->old.code_hdr->on_load_function_ptr == NULL) {
 	goto error;
     }
     if (BIF_ARG_2 != am_false && BIF_ARG_2 != am_true) {
@@ -667,44 +667,55 @@ 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.
+	 */
+	t = modp->curr;
+	modp->curr = modp->old;
+	modp->old = t;
 
 	/*
 	 * The on_load function succeded. Fix up export entries.
 	 */
 	for (i = 0; i < export_list_size(code_ix); i++) {
 	    Export *ep = export_list(i,code_ix);
-	    if (ep != NULL &&
-		ep->code[0] == BIF_ARG_1 &&
-		ep->code[4] != 0) {
+	    if (ep == NULL || ep->code[0] != BIF_ARG_1) {
+		continue;
+	    }
+	    if (ep->code[4] != 0) {
 		ep->addressv[code_ix] = (void *) ep->code[4];
 		ep->code[4] = 0;
+	    } else {
+		if (ep->addressv[code_ix] == ep->code+3 &&
+		    ep->code[3] == (BeamInstr) em_apply_bif) {
+		    continue;
+		}
+		ep->addressv[code_ix] = ep->code+3;
+		ep->code[3] = (BeamInstr) em_call_error_handler;
 	    }
 	}
 	modp->curr.code_hdr->on_load_function_ptr = NULL;
 	set_default_trace_pattern(BIF_ARG_1);
     } else if (BIF_ARG_2 == am_false) {
-	BeamInstr* code;
-	BeamInstr* end;
+	int i;
 
 	/*
-	 * The on_load function failed. Remove the loaded code.
-	 * This is an combination of delete and purge. We purge
-	 * the current code; the old code is not touched.
+	 * The on_load function failed. Remove references to the
+	 * code that is about to be purged from the export entries.
 	 */
-	erts_total_code_size -= modp->curr.code_length;
-	code = (BeamInstr*) modp->curr.code_hdr;
-	end = (BeamInstr *) ((char *)code + modp->curr.code_length);
-	erts_cleanup_funs_on_purge(code, end);
-	beam_catches_delmod(modp->curr.catches, code, modp->curr.code_length,
-			    erts_active_code_ix());
-        if (modp->curr.code_hdr->literals_start) {
-            erts_free(ERTS_ALC_T_LITERAL, modp->curr.code_hdr->literals_start);
-        }
-	erts_free(ERTS_ALC_T_CODE, modp->curr.code_hdr);
-	modp->curr.code_hdr = NULL;
-	modp->curr.code_length = 0;
-	modp->curr.catches = BEAM_CATCHES_NIL;
-	erts_remove_from_ranges(code);
+
+	for (i = 0; i < export_list_size(code_ix); i++) {
+	    Export *ep = export_list(i,code_ix);
+	    if (ep == NULL || ep->code[0] != BIF_ARG_1) {
+		continue;
+	    }
+	    if (ep->code[3] == (BeamInstr) em_apply_bif) {
+		continue;
+	    }
+	    ep->code[4] = 0;
+	}
     }
     erts_smp_thr_progress_unblock();
     erts_smp_proc_lock(BIF_P, ERTS_PROC_LOCK_MAIN);
diff --git a/erts/emulator/beam/beam_load.c b/erts/emulator/beam/beam_load.c
index 9bd9b4bfc0..95489d9a63 100644
--- a/erts/emulator/beam/beam_load.c
+++ b/erts/emulator/beam/beam_load.c
@@ -32,6 +32,7 @@
 #include "bif.h"
 #include "external.h"
 #include "beam_load.h"
+#include "beam_bp.h"
 #include "big.h"
 #include "erl_bits.h"
 #include "beam_catches.h"
@@ -766,7 +767,7 @@ Eterm
 erts_finish_loading(Binary* magic, Process* c_p,
 		    ErtsProcLocks c_p_locks, Eterm* modp)
 {
-    Eterm retval;
+    Eterm retval = NIL;
     LoaderState* stp = ERTS_MAGIC_BIN_DATA(magic);
     Module* mod_tab_p;
     struct erl_module_instance* inst_p;
@@ -779,16 +780,52 @@ erts_finish_loading(Binary* magic, Process* c_p,
 
     ERTS_SMP_LC_ASSERT(erts_initialized == 0 || erts_has_code_write_permission() ||
 		       erts_smp_thr_progress_is_blocking());
-
     /*
      * Make current code for the module old and insert the new code
      * as current.  This will fail if there already exists old code
      * for the module.
      */
 
+    mod_tab_p = erts_put_module(stp->module);
     CHKBLK(ERTS_ALC_T_CODE,stp->code);
-    retval = beam_make_current_old(c_p, c_p_locks, stp->module);
-    ASSERT(retval == NIL);
+    if (!stp->on_load) {
+	/*
+	 * Normal case -- no -on_load() function.
+	 */
+	retval = beam_make_current_old(c_p, c_p_locks, stp->module);
+	ASSERT(retval == NIL);
+    } else {
+	ErtsCodeIndex code_ix = erts_staging_code_ix();
+	Eterm module = stp->module;
+	int i;
+
+	/*
+	 * There is an -on_load() function. We will keep the current
+	 * code, but we must turn off any tracing.
+	 */
+
+	for (i = 0; i < export_list_size(code_ix); i++) {
+	    Export *ep = export_list(i, code_ix);
+	    if (ep == NULL || ep->code[0] != module) {
+		continue;
+	    }
+	    if (ep->addressv[code_ix] == ep->code+3) {
+		if (ep->code[3] == (BeamInstr) em_apply_bif) {
+		    continue;
+		} else if (ep->code[3] ==
+			   (BeamInstr) BeamOp(op_i_generic_breakpoint)) {
+		    ERTS_SMP_LC_ASSERT(erts_smp_thr_progress_is_blocking());
+		    ASSERT(mod_tab_p->curr.num_traced_exports > 0);
+		    erts_clear_export_break(mod_tab_p, ep->code+3);
+		    ep->addressv[code_ix] = (BeamInstr *) ep->code[4];
+		    ep->code[4] = 0;
+		}
+		ASSERT(ep->code[4] == 0);
+	    }
+	}
+	ASSERT(mod_tab_p->curr.num_breakpoints == 0);
+	ASSERT(mod_tab_p->curr.num_traced_exports == 0);
+    }
 
     /*
      * Update module table.
@@ -796,8 +833,11 @@ erts_finish_loading(Binary* magic, Process* c_p,
 
     size = stp->loaded_size;
     erts_total_code_size += size;
-    mod_tab_p = erts_put_module(stp->module);
-    inst_p = &mod_tab_p->curr;
+    if (stp->on_load) {
+	inst_p = &mod_tab_p->old;
+    } else {
+	inst_p = &mod_tab_p->curr;
+    }
     inst_p->code_hdr = stp->hdr;
     inst_p->code_length = size;
 
@@ -4757,10 +4797,10 @@ final_touch(LoaderState* stp, struct erl_module_instance* inst_p)
 	    ep->addressv[erts_staging_code_ix()] = address;
 	} else {
 	    /*
-	     * Don't make any of the exported functions
-	     * callable yet.
+	     * on_load: Don't make any of the exported functions
+	     * callable yet. Keep any function in the current
+	     * code callable.
 	     */
-	    ep->addressv[erts_staging_code_ix()] = ep->code+3;
 	    ep->code[4] = (BeamInstr) address;
 	}
     }
diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 941f44b9ec..dc83a780a2 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -2780,7 +2780,7 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
     ErlNifEntry* entry = NULL;
     ErlNifEnv env;
     int i, err, encoding;
-    Module* mod;
+    Module* module_p;
     Eterm mod_atom;
     const Atom* mod_atomp;
     Eterm f_atom;
@@ -2790,6 +2790,8 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
     int veto;
     struct erl_module_nif* lib = NULL;
     int reload_warning = 0;
+    struct erl_module_instance* this_mi;
+    struct erl_module_instance* prev_mi;
 
     encoding = erts_get_native_filename_encoding();
     if (encoding == ERL_FILENAME_WIN_WCHAR) {
@@ -2823,21 +2825,29 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
     ASSERT(caller != NULL);
     mod_atom = caller[0];
     ASSERT(is_atom(mod_atom));
-    mod=erts_get_module(mod_atom, erts_active_code_ix());
-    ASSERT(mod != NULL);
+    module_p = erts_get_module(mod_atom, erts_active_code_ix());
+    ASSERT(module_p != NULL);
 
     mod_atomp = atom_tab(atom_val(mod_atom));
     init_func = erts_static_nif_get_nif_init((char*)mod_atomp->name, mod_atomp->len);
     if (init_func != NULL)
       handle = init_func;
 
-    if (!in_area(caller, mod->curr.code_hdr, mod->curr.code_length)) {
-	ASSERT(in_area(caller, mod->old.code_hdr, mod->old.code_length));
+    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;
+	    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;
+	}
+    } else {
+	this_mi = &module_p->curr;
+	prev_mi = &module_p->old;
+    }
 
-	ret = load_nif_error(BIF_P, "old_code", "Calling load_nif from old "
-			     "module '%T' not allowed", mod_atom);
-    }    
-    else if (init_func == NULL &&
+    if (init_func == NULL &&
 	     (err=erts_sys_ddll_open(lib_name, &handle, &errdesc)) != ERL_DE_NO_ERROR) {
 	const char slogan[] = "Failed to load NIF library";
 	if (strstr(errdesc.str, lib_name) != NULL) {
@@ -2886,7 +2896,7 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 	for (i=0; i < entry->num_of_funcs && ret==am_ok; i++) {
 	    BeamInstr** code_pp;
 	    if (!erts_atom_get(f->name, sys_strlen(f->name), &f_atom, ERTS_ATOM_ENC_LATIN1)
-		|| (code_pp = get_func_pp(mod->curr.code_hdr, f_atom, f->arity))==NULL) {
+		|| (code_pp = get_func_pp(this_mi->code_hdr, f_atom, f->arity))==NULL) {
 		ret = load_nif_error(BIF_P,bad_lib,"Function not found %T:%s/%u",
 				     mod_atom, f->name, f->arity);
 	    }
@@ -2937,9 +2947,9 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
     erts_refc_init(&lib->rt_cnt, 0);
     erts_refc_init(&lib->rt_dtor_cnt, 0);
     ASSERT(opened_rt_list == NULL);
-    lib->mod = mod;
+    lib->mod = module_p;
     env.mod_nif = lib;
-    if (mod->curr.nif != NULL) { /*************** Reload ******************/
+    if (this_mi->nif != NULL) { /*************** Reload ******************/
 	/*
 	 * Repeated load_nif calls from same Erlang module instance ("reload")
 	 * is deprecated and was only ment as a development feature not to
@@ -2947,16 +2957,16 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 	 */
 	int k, old_incr = 0;
 	ErlNifFunc* old_func;
-        lib->priv_data = mod->curr.nif->priv_data;
+        lib->priv_data = this_mi->nif->priv_data;
 
-	ASSERT(mod->curr.nif->entry != NULL);
+	ASSERT(this_mi->nif->entry != NULL);
 	if (entry->reload == NULL) {
 	    ret = load_nif_error(BIF_P,reload,"Reload not supported by this NIF library.");
 	    goto error;
 	}
 	/* Check that no NIF is removed */
-	old_func = mod->curr.nif->entry->funcs;
-	for (k=0; k < mod->curr.nif->entry->num_of_funcs; k++) {
+	old_func = this_mi->nif->entry->funcs;
+	for (k=0; k < this_mi->nif->entry->num_of_funcs; k++) {
 	    int incr = 0;
 	    ErlNifFunc* f = entry->funcs;
 	    for (i=0; i < entry->num_of_funcs; i++) {
@@ -2972,7 +2982,7 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 				     old_func->name, old_func->arity);
 		goto error;
 	    }
-	    old_func = next_func(mod->curr.nif->entry, &old_incr, old_func);
+	    old_func = next_func(this_mi->nif->entry, &old_incr, old_func);
 	}
 	erts_pre_nif(&env, BIF_P, lib, NULL);
 	veto = entry->reload(&env, &lib->priv_data, BIF_ARG_2);
@@ -2982,24 +2992,24 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 	}
 	else {
 	    commit_opened_resource_types(lib);
-	    mod->curr.nif->entry = NULL; /* to prevent 'unload' callback */
-	    erts_unload_nif(mod->curr.nif);
+	    this_mi->nif->entry = NULL; /* to prevent 'unload' callback */
+	    erts_unload_nif(this_mi->nif);
 	    reload_warning = 1;
 	}
     }
     else {
 	lib->priv_data = NULL;
-	if (mod->old.nif != NULL) { /**************** Upgrade ***************/
-	    void* prev_old_data = mod->old.nif->priv_data;
+	if (prev_mi->nif != NULL) { /**************** Upgrade ***************/
+	    void* prev_old_data = prev_mi->nif->priv_data;
 	    if (entry->upgrade == NULL) {
 		ret = load_nif_error(BIF_P, upgrade, "Upgrade not supported by this NIF library.");
 		goto error;
 	    }
-	    erts_pre_nif(&env, BIF_P, lib, NULL);
-	    veto = entry->upgrade(&env, &lib->priv_data, &mod->old.nif->priv_data, BIF_ARG_2);
+		erts_pre_nif(&env, BIF_P, lib, NULL);
+	    veto = entry->upgrade(&env, &lib->priv_data, &prev_mi->nif->priv_data, BIF_ARG_2);
 	    erts_post_nif(&env);
 	    if (veto) {
-		mod->old.nif->priv_data = prev_old_data;
+		prev_mi->nif->priv_data = prev_old_data;
 		ret = load_nif_error(BIF_P, upgrade, "Library upgrade-call unsuccessful.");
 	    }
 	    else
@@ -3024,12 +3034,12 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 	int incr = 0;
 	ErlNifFunc* f = entry->funcs;
 
-	mod->curr.nif = lib;
+	this_mi->nif = lib;
 	for (i=0; i < entry->num_of_funcs; i++)
 	{
 	    BeamInstr* code_ptr;
 	    erts_atom_get(f->name, sys_strlen(f->name), &f_atom, ERTS_ATOM_ENC_LATIN1);
-	    code_ptr = *get_func_pp(mod->curr.code_hdr, f_atom, f->arity);
+	    code_ptr = *get_func_pp(this_mi->code_hdr, f_atom, f->arity);
 
 	    if (code_ptr[1] == 0) {
 		code_ptr[5+0] = (BeamInstr) BeamOp(op_call_nif);
diff --git a/lib/kernel/src/code_server.erl b/lib/kernel/src/code_server.erl
index 90b2a06c46..56372f90ca 100644
--- a/lib/kernel/src/code_server.erl
+++ b/lib/kernel/src/code_server.erl
@@ -1351,6 +1351,7 @@ finish_on_load_1(Mod, File, OnLoadRes, WaitingPids, Db) ->
     Res = case Keep of
 	      false ->
 		  _ = finish_on_load_report(Mod, OnLoadRes),
+		  _ = erts_code_purger:purge(Mod),
 		  {error,on_load_failure};
 	      true ->
 		  ets:insert(Db, {Mod,File}),
diff --git a/lib/kernel/test/code_SUITE.erl b/lib/kernel/test/code_SUITE.erl
index a0b72925f7..547881f3b2 100644
--- a/lib/kernel/test/code_SUITE.erl
+++ b/lib/kernel/test/code_SUITE.erl
@@ -20,6 +20,7 @@
 -module(code_SUITE).
 
 -include_lib("common_test/include/ct.hrl").
+-include_lib("syntax_tools/include/merl.hrl").
 
 -export([all/0, suite/0,groups/0,init_per_group/2,end_per_group/2]).
 -export([set_path/1, get_path/1, add_path/1, add_paths/1, del_path/1,
@@ -33,7 +34,9 @@
 	 where_is_file/1,
 	 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, big_boot_embedded/1,
+	 on_load_embedded/1, on_load_errors/1, on_load_update/1,
+	 on_load_purge/1,
+	 big_boot_embedded/1,
 	 native_early_modules/1, get_mode/1,
 	 normalized_paths/1]).
 
@@ -61,7 +64,8 @@ all() ->
      ext_mod_dep, clash, where_is_file,
      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_binary, on_load_embedded, on_load_errors, on_load_update,
+     on_load_purge,
      big_boot_embedded, native_early_modules, get_mode, normalized_paths].
 
 groups() ->
@@ -1436,6 +1440,89 @@ do_on_load_error(ReturnValue) ->
 	    {undef,[{on_load_error,main,[],_}|_]} = Exit
     end.
 
+on_load_update(_Config) ->
+    {Mod,Code1} = on_load_update_code(1),
+    {module,Mod} = code:load_binary(Mod, "", Code1),
+    42 = Mod:a(),
+    100 = Mod:b(99),
+    4 = erlang:trace_pattern({Mod,'_','_'}, true),
+
+    {Mod,Code2} = on_load_update_code(2),
+    {error,on_load_failure} = code:load_binary(Mod, "", Code2),
+    42 = Mod:a(),
+    100 = Mod:b(99),
+    {'EXIT',{undef,_}} = (catch Mod:never()),
+    4 = erlang:trace_pattern({Mod,'_','_'}, false),
+
+    {Mod,Code3} = on_load_update_code(3),
+    {module,Mod} = code:load_binary(Mod, "", Code3),
+    100 = Mod:c(),
+    {'EXIT',{undef,_}} = (catch Mod:a()),
+    {'EXIT',{undef,_}} = (catch Mod:b(10)),
+    {'EXIT',{undef,_}} = (catch Mod:never()),
+
+    ok.
+
+on_load_update_code(Version) ->
+    Mod = ?FUNCTION_NAME,
+    Tree = on_load_update_code_1(Version, Mod),
+    io:format("Version ~p", [Version]),
+    {ok,Mod,Code} = merl:compile(Tree),
+    merl:print(Tree),
+    io:nl(),
+    {Mod,Code}.
+
+on_load_update_code_1(1, Mod) ->
+    ?Q(["-module('@Mod@').\n",
+	"-export([a/0,b/1]).\n"
+	"-on_load(f/0).\n",
+	"f() -> ok.\n",
+	"a() -> 42.\n"
+	"b(I) -> I+1.\n"]);
+on_load_update_code_1(2, Mod) ->
+    ?Q(["-module('@Mod@').\n",
+	"-export([never/0]).\n"
+	"-on_load(f/0).\n",
+	"f() -> 42 = '@Mod@':a(), 1 = '@Mod@':b(0), fail.\n",
+	"never() -> never.\n"]);
+on_load_update_code_1(3, Mod) ->
+    ?Q(["-module('@Mod@').\n",
+	"-export([c/0]).\n"
+	"-on_load(f/0).\n",
+	"f() -> ok.\n",
+	"c() -> 100.\n"]).
+
+on_load_purge(_Config) ->
+    Mod = ?FUNCTION_NAME,
+    register(Mod, self()),
+    Tree = ?Q(["-module('@Mod@').\n",
+	       "-on_load(f/0).\n",
+	       "loop() -> loop().\n",
+	       "f() ->\n",
+	       "'@Mod@' ! {self(),spawn(fun loop/0)},\n",
+	       "receive Ack -> Ack end.\n"]),
+    merl:print(Tree),
+    {ok,Mod,Code} = merl:compile(Tree),
+    P = spawn(fun() ->
+		      exit(code:load_binary(Mod, "", Code))
+	      end),
+    monitor(process, P),
+    receive
+	{Pid1,Pid2} ->
+	    monitor(process, Pid2),
+	    Pid1 ! ack_and_failure,
+	    receive
+		{'DOWN',_,process,P,Exit1} ->
+		    {error,on_load_failure} = Exit1
+	    end,
+	    receive
+		{'DOWN',_,process,Pid2,Exit2} ->
+		    io:format("~p\n", [Exit2])
+	    after 10000 ->
+		    ct:fail(no_down_message)
+	    end
+    end.
+
 %% 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