Support for dirty BIFs

1 files changed, 253 insertions, 454 deletions

diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 5499512dd1..ee4b22f7ca 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -55,6 +55,9 @@
 #include "dtrace-wrapper.h"
 #include "erl_process.h"
 #include "erl_bif_unique.h"
+#undef ERTS_WANT_NFUNC_SCHED_INTERNALS__
+#define ERTS_WANT_NFUNC_SCHED_INTERNALS__
+#include "erl_nfunc_sched.h"
 #if defined(USE_DYNAMIC_TRACE) && (defined(USE_DTRACE) || defined(USE_SYSTEMTAP))
 #define HAVE_USE_DTRACE 1
 #endif
@@ -79,8 +82,11 @@ struct erl_module_nif {
     ErlNifFunc _funcs_copy_[1];  /* only used for old libs */
 };
 
+typedef ERL_NIF_TERM (*NativeFunPtr)(ErlNifEnv*, int, const ERL_NIF_TERM[]);
+
 #ifdef DEBUG
 #  define READONLY_CHECK
+#  define ERTS_DBG_NIF_NOT_SCHED_MARKER ((void *) (UWord) 1)
 #endif
 #ifdef READONLY_CHECK
 #  define ADD_READONLY_CHECK(ENV,PTR,SIZE) add_readonly_check(ENV,PTR,SIZE)
@@ -219,38 +225,6 @@ static void cache_env(ErlNifEnv* env);
 static void full_flush_env(ErlNifEnv *env);
 static void flush_env(ErlNifEnv* env);
 
-#ifdef ERTS_DIRTY_SCHEDULERS
-void erts_pre_dirty_nif(ErtsSchedulerData *esdp,
-			ErlNifEnv* env, Process* p,
-			struct erl_module_nif* mod_nif)
-{
-    Process *sproc;
-#ifdef DEBUG
-    erts_aint32_t state = erts_smp_atomic32_read_nob(&p->state);
-
-    ASSERT(!p->scheduler_data);
-    ASSERT((state & ERTS_PSFLG_DIRTY_RUNNING)
-	   && !(state & (ERTS_PSFLG_RUNNING|ERTS_PSFLG_RUNNING_SYS)));
-    ASSERT(esdp);
-#endif
-
-    erts_pre_nif(env, p, mod_nif, NULL);
-
-    sproc = esdp->dirty_shadow_process;
-    ASSERT(sproc);
-    ASSERT(sproc->static_flags & ERTS_STC_FLG_SHADOW_PROC);
-    ASSERT(erts_smp_atomic32_read_nob(&sproc->state)
-	   == (ERTS_PSFLG_ACTIVE
-	       | ERTS_PSFLG_DIRTY_RUNNING
-	       | ERTS_PSFLG_PROXY));
-
-    sproc->next = p;
-    sproc->common.id = p->common.id;
-    env->proc = sproc;
-    full_cache_env(env);
-}
-#endif
-
 /* Temporary object header, auto-deallocated when NIF returns
  * or when independent environment is cleared.
  */
@@ -278,115 +252,158 @@ void erts_post_nif(ErlNifEnv* env)
     env->exiting = ERTS_PROC_IS_EXITING(env->proc);
 }
 
-#ifdef ERTS_DIRTY_SCHEDULERS
-void erts_post_dirty_nif(ErlNifEnv* env)
+
+/*
+ * Initialize a NifExport struct. Create it if needed and store it in the
+ * proc. The direct_fp function is what will be invoked by op_call_nif, and
+ * the indirect_fp function, if not NULL, is what the direct_fp function
+ * will call. If the allocated NifExport isn't enough to hold all of argv,
+ * allocate a larger one. Save 'current' and registers if first time this
+ * call is scheduled.
+ */
+
+static ERTS_INLINE ERL_NIF_TERM
+schedule(ErlNifEnv* env, NativeFunPtr direct_fp, NativeFunPtr indirect_fp,
+	 Eterm mod, Eterm func_name, int argc, const ERL_NIF_TERM argv[])
 {
-    Process *c_p;
-    ASSERT(env->proc->static_flags & ERTS_STC_FLG_SHADOW_PROC);
-    ASSERT(env->proc->next);
-    erts_unblock_fpe(env->fpe_was_unmasked);
-    full_flush_env(env);
-    free_tmp_objs(env);
-    c_p = env->proc->next;
-    env->exiting = ERTS_PROC_IS_EXITING(c_p);
-    ERTS_VBUMP_ALL_REDS(c_p);
+    Export *exp;
+    NifExport *ep;
+    Process *c_p, *dirty_shadow_proc;
+
+    execution_state(env, &c_p, NULL);
+    if (c_p == env->proc)
+	dirty_shadow_proc = NULL;
+    else
+	dirty_shadow_proc = env->proc;
+
+    ERTS_SMP_LC_ASSERT(ERTS_PROC_LOCK_MAIN & erts_proc_lc_my_proc_locks(c_p));
+
+    exp = ErtsContainerStruct(c_p->current, Export, info.mfa);
+
+    ep = erts_nif_export_schedule(c_p, dirty_shadow_proc,
+				  c_p->current,
+				  (BifFunction) exp->beam[1],
+				  c_p->cp,
+				  (BeamInstr) em_call_nif,
+				  direct_fp, indirect_fp,
+				  mod, func_name,
+				  argc, (const Eterm *) argv);
+    if (!ep->m) {
+	/* First time this call is scheduled... */
+	erts_refc_inc(&env->mod_nif->rt_dtor_cnt, 1);
+	ep->m = env->mod_nif;
+    }
+    return (ERL_NIF_TERM) THE_NON_VALUE;
 }
-#endif
 
-static void full_flush_env(ErlNifEnv* env)
-{
 #ifdef ERTS_DIRTY_SCHEDULERS
-    if (env->proc->static_flags & ERTS_STC_FLG_SHADOW_PROC) {
-	/* Dirty nif call using shadow process struct */
-	Process *c_p = env->proc->next;
-
-	ASSERT(is_scheduler() < 0);
-	ASSERT(env->proc->common.id == c_p->common.id);
-	ERTS_SMP_LC_ASSERT(erts_proc_lc_my_proc_locks(c_p)
-			   & ERTS_PROC_LOCK_MAIN);
-
-	if (!env->heap_frag) {
-	    ASSERT(env->hp_end == HEAP_LIMIT(c_p));
-	    ASSERT(env->hp >= HEAP_TOP(c_p));
-	    ASSERT(env->hp <= HEAP_LIMIT(c_p));
-	    HEAP_TOP(c_p) = env->hp;
+
+static ERL_NIF_TERM dirty_nif_finalizer(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
+static ERL_NIF_TERM dirty_nif_exception(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
+
+int
+erts_call_dirty_nif(ErtsSchedulerData *esdp, Process *c_p, BeamInstr *I, Eterm *reg)
+{
+    int exiting;
+    ERL_NIF_TERM *argv = (ERL_NIF_TERM *) reg;
+    NifExport *nep = ERTS_I_BEAM_OP_TO_NIF_EXPORT(I);
+    ErtsCodeMFA *codemfa = erts_code_to_codemfa(I);
+    NativeFunPtr dirty_nif = (NativeFunPtr) I[1];
+    ErlNifEnv env;
+    ERL_NIF_TERM result;
+#ifdef DEBUG
+    erts_aint32_t state = erts_smp_atomic32_read_nob(&c_p->state);
+
+    ASSERT(nep == ERTS_PROC_GET_NIF_TRAP_EXPORT(c_p));
+
+    ASSERT(!c_p->scheduler_data);
+    ASSERT((state & ERTS_PSFLG_DIRTY_RUNNING)
+	&& !(state & (ERTS_PSFLG_RUNNING|ERTS_PSFLG_RUNNING_SYS)));
+    ASSERT(esdp);
+
+    nep->func = ERTS_DBG_NIF_NOT_SCHED_MARKER;
+#endif
+
+    erts_pre_nif(&env, c_p, nep->m, NULL);
+
+    env.proc = erts_make_dirty_shadow_proc(esdp, c_p);
+
+    env.proc->freason = EXC_NULL;
+    env.proc->fvalue = NIL;
+    env.proc->ftrace = NIL;
+    env.proc->i = c_p->i;
+
+    ASSERT(ERTS_SCHEDULER_IS_DIRTY(erts_proc_sched_data(c_p)));
+
+    erts_smp_atomic32_read_band_mb(&c_p->state, ~(ERTS_PSFLG_DIRTY_CPU_PROC
+						   | ERTS_PSFLG_DIRTY_IO_PROC));
+
+    erts_smp_proc_unlock(c_p, ERTS_PROC_LOCK_MAIN);
+
+    result = (*dirty_nif)(&env, codemfa->arity, argv); /* Call dirty NIF */
+
+    erts_smp_proc_lock(c_p, ERTS_PROC_LOCK_MAIN);
+
+    ASSERT(env.proc->static_flags & ERTS_STC_FLG_SHADOW_PROC);
+    ASSERT(env.proc->next == c_p);
+
+    exiting = ERTS_PROC_IS_EXITING(c_p);
+
+    if (!exiting) {
+	if (env.exception_thrown) {
+	schedule_exception:
+	    schedule(&env, dirty_nif_exception, NULL,
+		     am_erts_internal, am_dirty_nif_exception,
+		     1, &env.proc->fvalue);
+	}
+	else if (is_value(result)) {
+	    schedule(&env, dirty_nif_finalizer, NULL,
+		     am_erts_internal, am_dirty_nif_finalizer,
+		     1, &result);
+	}
+	else if (env.proc->freason != TRAP) { /* user returned garbage... */
+	    ERTS_DECL_AM(badreturn);
+	    (void) enif_raise_exception(&env, AM_badreturn);
+	    goto schedule_exception;
 	}
 	else {
-	    Uint usz;
-	    ASSERT(env->hp_end != HEAP_LIMIT(c_p));
-	    ASSERT(env->hp_end - env->hp <= env->heap_frag->alloc_size);
-
-	    HEAP_TOP(c_p) = HEAP_TOP(env->proc);
-	    usz = env->hp - env->heap_frag->mem;
-	    env->proc->mbuf_sz += usz - env->heap_frag->used_size;
-	    env->heap_frag->used_size = usz;
-
-	    ASSERT(env->heap_frag->used_size <= env->heap_frag->alloc_size);
-
-	    if (c_p->mbuf) {
-		ErlHeapFragment *bp;
-		for (bp = env->proc->mbuf; bp->next; bp = bp->next)
-		    ;
-		bp->next = c_p->mbuf;
-	    }
+	    /* Rescheduled by dirty NIF call... */
+	    ASSERT(nep->func != ERTS_DBG_NIF_NOT_SCHED_MARKER);
+	}
+	c_p->i = env.proc->i;
+	c_p->arity = env.proc->arity;
+    }
 
-	    c_p->mbuf = env->proc->mbuf;
-	    c_p->mbuf_sz += env->proc->mbuf_sz;
+#ifdef DEBUG
+    if (nep->func == ERTS_DBG_NIF_NOT_SCHED_MARKER)
+	nep->func = NULL;
+#endif
 
-	}
+    erts_unblock_fpe(env.fpe_was_unmasked);
+    full_flush_env(&env);
+    free_tmp_objs(&env);
 
-	if (!c_p->off_heap.first)
-	    c_p->off_heap.first = env->proc->off_heap.first;
-	else if (env->proc->off_heap.first) {
-	    struct erl_off_heap_header *ohhp;
-	    for (ohhp = env->proc->off_heap.first; ohhp->next; ohhp = ohhp->next)
-		;
-	    ohhp->next = c_p->off_heap.first;
-	    c_p->off_heap.first = env->proc->off_heap.first;
-	}
-	c_p->off_heap.overhead += env->proc->off_heap.overhead;
+    return exiting;
+}
 
-	return;
-    }
 #endif
 
+static void full_flush_env(ErlNifEnv* env)
+{
     flush_env(env);
+#ifdef ERTS_DIRTY_SCHEDULERS
+    if (env->proc->static_flags & ERTS_STC_FLG_SHADOW_PROC)
+	/* Dirty nif call using shadow process struct */
+	erts_flush_dirty_shadow_proc(env->proc);
+#endif
 }
 
 static void full_cache_env(ErlNifEnv* env)
 {    
 #ifdef ERTS_DIRTY_SCHEDULERS
-    if (env->proc->static_flags & ERTS_STC_FLG_SHADOW_PROC) {
-	/* Dirty nif call using shadow process struct */
-	Process *sproc = env->proc;
-	Process *c_p = sproc->next;
-	ASSERT(c_p);
-	ASSERT(is_scheduler() < 0);
-	ASSERT(env->proc->common.id == c_p->common.id);
-	ERTS_SMP_LC_ASSERT(erts_proc_lc_my_proc_locks(c_p)
-			   & ERTS_PROC_LOCK_MAIN);
-
-	sproc->htop = c_p->htop;
-	sproc->stop = c_p->stop;
-	sproc->hend = c_p->hend;
-	sproc->heap = c_p->heap;
-	sproc->abandoned_heap = c_p->abandoned_heap;
-	sproc->heap_sz = c_p->heap_sz;
-	sproc->high_water = c_p->high_water;
-	sproc->old_hend = c_p->old_hend;
-	sproc->old_htop = c_p->old_htop;
-	sproc->old_heap = c_p->old_heap;
-	sproc->mbuf = NULL;
-	sproc->mbuf_sz = 0;
-	ERTS_INIT_OFF_HEAP(&sproc->off_heap);
-
-	env->hp_end = HEAP_LIMIT(c_p);
-	env->hp = HEAP_TOP(c_p);
-	env->heap_frag = NULL;
-	return;
-    }
+    if (env->proc->static_flags & ERTS_STC_FLG_SHADOW_PROC)
+	erts_cache_dirty_shadow_proc(env->proc);
 #endif
-
     cache_env(env);
 }
 
@@ -1303,22 +1320,15 @@ Eterm enif_make_badarg(ErlNifEnv* env)
 
 Eterm enif_raise_exception(ErlNifEnv* env, ERL_NIF_TERM reason)
 {
-    Process *c_p;
-
-    execution_state(env, &c_p, NULL);
-
     env->exception_thrown = 1;
-    c_p->fvalue = reason;
-    BIF_ERROR(c_p, EXC_ERROR);
+    env->proc->fvalue = reason;
+    BIF_ERROR(env->proc, EXC_ERROR);
 }
 
 int enif_has_pending_exception(ErlNifEnv* env, ERL_NIF_TERM* reason)
 {
-    if (env->exception_thrown && reason != NULL) {
-	Process *c_p;
-	execution_state(env, &c_p, NULL);
-	*reason = c_p->fvalue;
-    }
+    if (env->exception_thrown && reason != NULL)
+	*reason = env->proc->fvalue;
     return env->exception_thrown;
 }
 
@@ -2269,188 +2279,28 @@ int enif_consume_timeslice(ErlNifEnv* env, int percent)
     return ERTS_BIF_REDS_LEFT(proc) == 0;
 }
 
-/*
- * NIF exports need a few more items than the Export struct provides,
- * including the erl_module_nif* and a NIF function pointer, so the
- * NifExport below adds those. The Export member must be first in the
- * struct. The saved_current, exception_thrown, saved_argc, rootset_extra, and
- * rootset members are used to track the MFA, any pending exception, and
- * arguments of the top NIF in case a chain of one or more
- * enif_schedule_nif() calls results in an exception, since in that case
- * the original MFA and registers have to be restored before returning to
- * Erlang to ensure stacktrace information associated with the exception is
- * correct.
- */
-typedef ERL_NIF_TERM (*NativeFunPtr)(ErlNifEnv*, int, const ERL_NIF_TERM[]);
-
-typedef struct {
-    Export exp;
-    struct erl_module_nif* m;
-    NativeFunPtr fp;
-    ErtsCodeMFA *saved_current;
-    int exception_thrown;
-    int saved_argc;
-    int rootset_extra;
-    Eterm rootset[1];
-} NifExport;
-
-/*
- * If a process has saved arguments, they need to be part of the GC
- * rootset. The function below is called from setup_rootset() in
- * erl_gc.c. This function is declared in erl_process.h. Any exception term
- * saved in the NifExport is also made part of the GC rootset here; it
- * always resides in rootset[0].
- */
-int
-erts_setup_nif_gc(Process* proc, Eterm** objv, int* nobj)
-{
-    NifExport* ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    int gc = ep && (ep->saved_argc > 0 || ep->rootset[0] != NIL);
-
-    if (gc) {
-	*objv = ep->rootset;
-	*nobj = 1 + ep->saved_argc;
-    }
-    return gc;
-}
-
-int
-erts_check_nif_export_in_area(Process *p, char *start, Uint size)
-{
-    NifExport *nep = ERTS_PROC_GET_NIF_TRAP_EXPORT(p);
-    if (!nep || !nep->saved_current)
-	return 0;
-    if (ErtsInArea(nep->saved_current, start, size))
-	return 1;
-    return 0;
-}
-
-/*
- * Allocate a NifExport and set it in proc specific data
- */
-static NifExport*
-allocate_nif_sched_data(Process* proc, int argc)
-{
-    NifExport* ep;
-    size_t total;
-    int i;
-
-    total = sizeof(NifExport) + argc*sizeof(Eterm);
-    ep = erts_alloc(ERTS_ALC_T_NIF_TRAP_EXPORT, total);
-    sys_memset((void*) ep, 0, total);
-    ep->rootset_extra = argc;
-    ep->rootset[0] = NIL;
-    for (i=0; i<ERTS_NUM_CODE_IX; i++) {
-	ep->exp.addressv[i] = &ep->exp.beam[0];
-    }
-    ep->exp.beam[0] = (BeamInstr) em_call_nif;
-    (void) ERTS_PROC_SET_NIF_TRAP_EXPORT(proc, ep);
-    return ep;
-}
-
 static ERTS_INLINE void
-destroy_nif_export(NifExport *nif_export)
+nif_export_cleanup_nif_mod(NifExport *ep)
 {
-    erts_free(ERTS_ALC_T_NIF_TRAP_EXPORT, (void *) nif_export);
+    if (erts_refc_dectest(&ep->m->rt_dtor_cnt, 0) == 0 && ep->m->mod == NULL)
+	close_lib(ep->m);
+    ep->m = NULL;
 }
 
 void
-erts_destroy_nif_export(void *nif_export)
+erts_nif_export_cleanup_nif_mod(NifExport *ep)
 {
-    destroy_nif_export((NifExport *) nif_export);
+    nif_export_cleanup_nif_mod(ep);
 }
 
-/*
- * Initialize a NifExport struct. Create it if needed and store it in the
- * proc. The direct_fp function is what will be invoked by op_call_nif, and
- * the indirect_fp function, if not NULL, is what the direct_fp function
- * will call. If the allocated NifExport isn't enough to hold all of argv,
- * allocate a larger one. Save MFA and registers only if the need_save
- * parameter is true.
- */
-static ERL_NIF_TERM
-init_nif_sched_data(ErlNifEnv* env, NativeFunPtr direct_fp, NativeFunPtr indirect_fp,
-		    int need_save, int argc, const ERL_NIF_TERM argv[])
+static ERTS_INLINE void
+nif_export_restore(Process *c_p, NifExport *ep)
 {
-    Process* proc;
-    Eterm* reg;
-    NifExport* ep;
-    int i, scheduler;
-    int orig_argc;
-
-    execution_state(env, &proc, &scheduler);
-
-    ASSERT(scheduler);
-
-    ERTS_SMP_LC_ASSERT(erts_proc_lc_my_proc_locks(proc)
-		       & ERTS_PROC_LOCK_MAIN);
-
-    reg = erts_proc_sched_data(proc)->x_reg_array;
-
-    ASSERT(!need_save || proc->current);
-    orig_argc = need_save ? (int) proc->current->arity : 0;
-
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    if (!ep)
-	ep = allocate_nif_sched_data(proc, orig_argc);
-    else if (need_save && ep->rootset_extra < orig_argc) {
-	NifExport* new_ep = allocate_nif_sched_data(proc, orig_argc);
-	destroy_nif_export(ep);
-	ep = new_ep;
-    }
-    if (env->exception_thrown) {
-	ep->exception_thrown = 1;
-	ep->rootset[0] = proc->fvalue;
-    } else {
-	ep->exception_thrown = 0;
-	ep->rootset[0] = NIL;
-    }
-    if (scheduler > 0)
-	ERTS_VBUMP_ALL_REDS(proc);
-    if (need_save) {
-	ep->saved_current = proc->current;
-	ep->saved_argc = orig_argc;
-	for (i = 0; i < orig_argc; i++)
-	    ep->rootset[i+1] = reg[i];
-    }
-    for (i = 0; i < argc; i++)
-	reg[i] = (Eterm) argv[i];
-    proc->i = (BeamInstr*) ep->exp.addressv[0];
-    ep->exp.info.mfa.module = proc->current->module;
-    ep->exp.info.mfa.function = proc->current->function;
-    ep->exp.info.mfa.arity = argc;
-    ep->exp.beam[1] = (BeamInstr) direct_fp;
-    ep->m = env->mod_nif;
-    ep->fp = indirect_fp;
-    proc->freason = TRAP;
-    proc->arity = argc;
-    return THE_NON_VALUE;
+    erts_nif_export_restore(c_p, ep);
+    ASSERT(ep->m);
+    nif_export_cleanup_nif_mod(ep);
 }
 
-/*
- * Restore saved MFA and registers. Registers are restored only when the
- * exception flag is true.
- */
-static void
-restore_nif_mfa(Process* proc, NifExport* ep, int exception)
-{
-    int i;
-
-    ERTS_SMP_LC_ASSERT(!(proc->static_flags
-			 & ERTS_STC_FLG_SHADOW_PROC));
-    ERTS_SMP_LC_ASSERT(erts_proc_lc_my_proc_locks(proc)
-		       & ERTS_PROC_LOCK_MAIN);
-
-    ASSERT(ep->saved_current != &ep->exp.info.mfa);
-    proc->current = ep->saved_current;
-    ep->saved_current = NULL;
-    if (exception) {
-	Eterm* reg = erts_proc_sched_data(proc)->x_reg_array;
-	for (i = 0; i < ep->saved_argc; i++)
-	    reg[i] = ep->rootset[i+1];
-    }
-    ep->saved_argc = 0;
-}
 
 #ifdef ERTS_DIRTY_SCHEDULERS
 
@@ -2459,7 +2309,7 @@ restore_nif_mfa(Process* proc, NifExport* ep, int exception)
  * switch the process off a dirty scheduler thread and back onto a regular
  * scheduler thread, and then return the result from the dirty NIF. It also
  * restores the original NIF MFA when necessary based on the value of
- * ep->fp set by execute_dirty_nif via init_nif_sched_data -- non-NULL
+ * ep->func set by execute_dirty_nif via init_nif_sched_data -- non-NULL
  * means restore, NULL means do not restore.
  */
 static ERL_NIF_TERM
@@ -2474,9 +2324,7 @@ dirty_nif_finalizer(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     ASSERT(!ERTS_SCHEDULER_IS_DIRTY(erts_proc_sched_data(proc)));
     ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
     ASSERT(ep);
-    ASSERT(!ep->exception_thrown);
-    if (ep->fp)
-	restore_nif_mfa(proc, ep, 0);
+    nif_export_restore(proc, ep);
     return argv[0];
 }
 
@@ -2486,148 +2334,100 @@ dirty_nif_finalizer(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 static ERL_NIF_TERM
 dirty_nif_exception(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {
+    ERL_NIF_TERM ret;
     Process* proc;
     NifExport* ep;
+    Eterm exception;
 
     execution_state(env, &proc, NULL);
 
+    ASSERT(argc == 1);
     ASSERT(!ERTS_SCHEDULER_IS_DIRTY(erts_proc_sched_data(proc)));
     ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
     ASSERT(ep);
-    ASSERT(ep->exception_thrown);
-    if (ep->fp)
-	restore_nif_mfa(proc, ep, 1);
-    return enif_raise_exception(env, ep->rootset[0]);
+    exception = argv[0]; /* argv overwritten by restore below... */
+    nif_export_cleanup_nif_mod(ep);
+    ret = enif_raise_exception(env, exception);
+
+    /* Restore orig info for error and clear nif export in handle_error() */
+    proc->freason |= EXF_RESTORE_NIF;
+    return ret;
 }
 
 /*
- * Dirty NIF execution wrapper function. Invoke an application's dirty NIF,
- * then check the result and schedule the appropriate finalizer function
- * where needed. Also restore the original NIF MFA when appropriate.
+ * Dirty NIF scheduling wrapper function. Schedule a dirty NIF to execute.
+ * The dirty scheduler thread type (CPU or I/O) is indicated in flags
+ * parameter.
  */
-static ERL_NIF_TERM
-execute_dirty_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+static ERTS_INLINE ERL_NIF_TERM
+schedule_dirty_nif(ErlNifEnv* env, int flags, NativeFunPtr fp,
+		   Eterm func_name, int argc, const ERL_NIF_TERM argv[])
 {
     Process* proc;
-    NativeFunPtr fp;
-    NifExport* ep;
-    ERL_NIF_TERM result;
-
-    execution_state(env, &proc, NULL);
 
-    ep = ErtsContainerStruct(proc->current, NifExport, exp.info.mfa);
-    fp = ep->fp;
-
-    ASSERT(ERTS_SCHEDULER_IS_DIRTY(erts_proc_sched_data(proc)));
-
-    /*
-     * Set ep->fp to NULL before the native call so we know later whether it scheduled another NIF for execution
-     */
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    ASSERT(ep && fp);
-
-    ep->fp = NULL;
-    erts_smp_atomic32_read_band_mb(&proc->state, ~(ERTS_PSFLG_DIRTY_CPU_PROC
-						   | ERTS_PSFLG_DIRTY_IO_PROC));
+    ASSERT(is_atom(func_name));
+    ASSERT(fp);
 
-    erts_smp_proc_unlock(proc, ERTS_PROC_LOCK_MAIN);
+    ASSERT(flags==ERL_NIF_DIRTY_JOB_IO_BOUND || flags==ERL_NIF_DIRTY_JOB_CPU_BOUND);
 
-    result = (*fp)(env, argc, argv);
+    execution_state(env, &proc, NULL);
 
-    erts_smp_proc_lock(proc, ERTS_PROC_LOCK_MAIN);
+    (void) erts_smp_atomic32_read_bset_nob(&proc->state,
+					   (ERTS_PSFLG_DIRTY_CPU_PROC
+					    | ERTS_PSFLG_DIRTY_IO_PROC),
+					   (flags == ERL_NIF_DIRTY_JOB_CPU_BOUND
+					    ? ERTS_PSFLG_DIRTY_CPU_PROC
+					    : ERTS_PSFLG_DIRTY_IO_PROC));
 
-    if (erts_refc_dectest(&env->mod_nif->rt_dtor_cnt, 0) == 0 && env->mod_nif->mod == NULL)
-	close_lib(env->mod_nif);
-    /*
-     * If no more NIFs were scheduled by the native call via
-     * enif_schedule_nif(), then ep->fp will still be NULL as set above, in
-     * which case we need to restore the original NIF calling
-     * context. Reuse fp essentially as a boolean for this, passing it to
-     * init_nif_sched_data below. Both dirty_nif_exception and
-     * dirty_nif_finalizer then check ep->fp to decide whether or not to
-     * restore the original calling context.
-     */
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    ASSERT(ep);
-    if (ep->fp)
-	fp = NULL;
-    if (is_non_value(result) || env->exception_thrown) {
-	if (proc->freason != TRAP) {
-	    return init_nif_sched_data(env, dirty_nif_exception, fp, 0, argc, argv);
-	} else {
-	    if (ep->fp == NULL)
-		restore_nif_mfa(proc, ep, 1);
-	    return THE_NON_VALUE;
-	}
-    }
-    else
-	return init_nif_sched_data(env, dirty_nif_finalizer, fp, 0, 1, &result);
+    return schedule(env, fp, NULL, proc->current->module, func_name, argc, argv);
 }
 
-/*
- * Dirty NIF scheduling wrapper function. Schedule a dirty NIF to execute
- * via the execute_dirty_nif() wrapper function. The dirty scheduler thread
- * type (CPU or I/O) is indicated in flags parameter.
- */
 static ERTS_INLINE ERL_NIF_TERM
-schedule_dirty_nif(ErlNifEnv* env, int flags, int argc, const ERL_NIF_TERM argv[])
+static_schedule_dirty_nif(ErlNifEnv* env, erts_aint32_t dirty_psflg,
+			     int argc, const ERL_NIF_TERM argv[])
 {
-    ERL_NIF_TERM result;
-    erts_aint32_t act, dirty_flag;
-    Process* proc;
+    Process *proc;
+    NifExport *ep;
+    Eterm mod, func;
     NativeFunPtr fp;
-    NifExport* ep;
-    int need_save, scheduler;
 
-    execution_state(env, &proc, &scheduler);
-    if (scheduler <= 0) {
-	ASSERT(scheduler < 0);
-	erts_smp_proc_lock(proc, ERTS_PROC_LOCK_MAIN);
-    }
+    execution_state(env, &proc, NULL);
+
+    /*
+     * Called in order to schedule statically determined
+     * dirty NIF calls...
+     *
+     * Note that 'current' does not point into a NifExport
+     * structure; only a structure with similar
+     * parts (located in code).
+     */
 
     ep = ErtsContainerStruct(proc->current, NifExport, exp.info.mfa);
-    fp = ep->fp;
+    mod = proc->current->module;
+    func = proc->current->function;
+    fp = (NativeFunPtr) ep->func;
 
+    ASSERT(is_atom(mod) && is_atom(func));
     ASSERT(fp);
 
-    ASSERT(flags==ERL_NIF_DIRTY_JOB_IO_BOUND || flags==ERL_NIF_DIRTY_JOB_CPU_BOUND);
-
-    if (flags == ERL_NIF_DIRTY_JOB_CPU_BOUND)
-	dirty_flag = ERTS_PSFLG_DIRTY_CPU_PROC;
-    else
-	dirty_flag = ERTS_PSFLG_DIRTY_IO_PROC;
-
-    act = erts_smp_atomic32_read_bor_nob(&proc->state, dirty_flag);
-    if (!(act & (ERTS_PSFLG_DIRTY_CPU_PROC|ERTS_PSFLG_DIRTY_IO_PROC)))
-	erts_refc_inc(&env->mod_nif->rt_dtor_cnt, 1);
-    else if ((act & (ERTS_PSFLG_DIRTY_CPU_PROC
-		     | ERTS_PSFLG_DIRTY_IO_PROC)) & ~dirty_flag) {
-	/* clear other flag... */
-	if (flags == ERL_NIF_DIRTY_JOB_CPU_BOUND)
-	    dirty_flag = ERTS_PSFLG_DIRTY_IO_PROC;
-	else
-	    dirty_flag = ERTS_PSFLG_DIRTY_CPU_PROC;
-	erts_smp_atomic32_read_band_nob(&proc->state, ~dirty_flag);
-    }
+    (void) erts_smp_atomic32_read_bset_nob(&proc->state,
+					   (ERTS_PSFLG_DIRTY_CPU_PROC
+					    | ERTS_PSFLG_DIRTY_IO_PROC),
+					   dirty_psflg);
 
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    need_save = (ep == NULL || !ep->saved_current);
-    result = init_nif_sched_data(env, execute_dirty_nif, fp, need_save, argc, argv);
-    if (scheduler <= 0)
-	erts_smp_proc_unlock(proc, ERTS_PROC_LOCK_MAIN);
-    return result;
+    return schedule(env, fp, NULL, mod, func, argc, argv);
 }
 
 static ERL_NIF_TERM
-schedule_dirty_io_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+static_schedule_dirty_io_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {
-    return schedule_dirty_nif(env, ERL_NIF_DIRTY_JOB_IO_BOUND, argc, argv);
+    return static_schedule_dirty_nif(env, ERTS_PSFLG_DIRTY_IO_PROC, argc, argv);
 }
 
 static ERL_NIF_TERM
-schedule_dirty_cpu_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+static_schedule_dirty_cpu_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
 {
-    return schedule_dirty_nif(env, ERL_NIF_DIRTY_JOB_CPU_BOUND, argc, argv);
+    return static_schedule_dirty_nif(env, ERTS_PSFLG_DIRTY_CPU_PROC, argc, argv);
 }
 
 #endif /* ERTS_DIRTY_SCHEDULERS */
@@ -2646,23 +2446,42 @@ execute_nif(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     ERL_NIF_TERM result;
 
     execution_state(env, &proc, NULL);
-    ep = ErtsContainerStruct(proc->current, NifExport, exp.info.mfa);
-    fp = ep->fp;
 
-    ASSERT(!env->exception_thrown);
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
+    ep = ErtsContainerStruct(proc->current, NifExport, exp.info.mfa);
+    fp = ep->func;
     ASSERT(ep);
-    ep->fp = NULL;
+    ASSERT(!env->exception_thrown);
+
+    fp = (NativeFunPtr) ep->func;
+
+#ifdef DEBUG
+    ep->func = ERTS_DBG_NIF_NOT_SCHED_MARKER;
+#endif
+
     result = (*fp)(env, argc, argv);
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    ASSERT(ep);
-    /*
-     * If no NIFs were scheduled by the native call via
-     * enif_schedule_nif(), then ep->fp will still be NULL as set above, in
-     * which case we need to restore the original NIF MFA.
-     */
-    if (ep->fp == NULL)
-	restore_nif_mfa(proc, ep, env->exception_thrown);
+
+    ASSERT(ep == ERTS_PROC_GET_NIF_TRAP_EXPORT(proc));
+
+    if (is_value(result) || proc->freason != TRAP) {
+	/* Done (not rescheduled)... */
+	ASSERT(ep->func == ERTS_DBG_NIF_NOT_SCHED_MARKER);
+	if (!env->exception_thrown)
+	    nif_export_restore(proc, ep);
+	else {
+	    nif_export_cleanup_nif_mod(ep);
+	    /*
+	     * Restore orig info for error and clear nif
+	     * export in handle_error()
+	     */
+	    proc->freason |= EXF_RESTORE_NIF;
+	}
+    }
+
+#ifdef DEBUG
+    if (ep->func == ERTS_DBG_NIF_NOT_SCHED_MARKER)
+	ep->func = NULL;
+#endif
+
     return result;
 }
 
@@ -2672,9 +2491,8 @@ enif_schedule_nif(ErlNifEnv* env, const char* fun_name, int flags,
 		  int argc, const ERL_NIF_TERM argv[])
 {
     Process* proc;
-    NifExport* ep;
     ERL_NIF_TERM fun_name_atom, result;
-    int need_save, scheduler;
+    int scheduler;
 
     if (argc > MAX_ARG)
 	return enif_make_badarg(env);
@@ -2689,35 +2507,16 @@ enif_schedule_nif(ErlNifEnv* env, const char* fun_name, int flags,
 	erts_smp_proc_lock(proc, ERTS_PROC_LOCK_MAIN);
     }
 
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    need_save = (ep == NULL || !ep->saved_current);
-
-    if (flags) {
+    if (flags == 0)
+	result = schedule(env, execute_nif, fp, proc->current->module,
+			  fun_name_atom, argc, argv);
 #ifdef ERTS_DIRTY_SCHEDULERS
-	NativeFunPtr sched_fun;
-	int chkflgs = (flags & (ERL_NIF_DIRTY_JOB_IO_BOUND|ERL_NIF_DIRTY_JOB_CPU_BOUND));
-	if (chkflgs == ERL_NIF_DIRTY_JOB_IO_BOUND)
-	    sched_fun = schedule_dirty_io_nif;
-	else if (chkflgs == ERL_NIF_DIRTY_JOB_CPU_BOUND)
-	    sched_fun = schedule_dirty_cpu_nif;
-	else {
-	    result = enif_make_badarg(env);
-	    goto done;
-	}
-	result = init_nif_sched_data(env, sched_fun, fp, need_save, argc, argv);
-#else
-	result = enif_make_badarg(env);
+    else if (!(flags & ~(ERL_NIF_DIRTY_JOB_IO_BOUND|ERL_NIF_DIRTY_JOB_CPU_BOUND)))
+	result = schedule_dirty_nif(env, flags, fp, fun_name_atom, argc, argv);
 #endif
-	goto done;
-    }
     else
-	result = init_nif_sched_data(env, execute_nif, fp, need_save, argc, argv);
-
-    ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
-    ASSERT(ep);
-    ep->exp.info.mfa.function = (BeamInstr) fun_name_atom;
+	result = enif_make_badarg(env);
 
-done:
     if (scheduler < 0)
 	erts_smp_proc_unlock(proc, ERTS_PROC_LOCK_MAIN);
 
@@ -3416,8 +3215,8 @@ BIF_RETTYPE load_nif_2(BIF_ALIST_2)
 	    if (f->flags) {
 		code_ptr[3] = (BeamInstr) f->fptr;
 		code_ptr[1] = (f->flags == ERL_NIF_DIRTY_JOB_IO_BOUND) ?
-		    (BeamInstr) schedule_dirty_io_nif :
-		    (BeamInstr) schedule_dirty_cpu_nif;
+		    (BeamInstr) static_schedule_dirty_io_nif :
+		    (BeamInstr) static_schedule_dirty_cpu_nif;
 	    }
 	    else
 #endif