5 files changed, 21 insertions, 0 deletions
diff --git a/erts/emulator/beam/bif_instrs.tab b/erts/emulator/beam/bif_instrs.tab
index 0932b8b985..0f074280db 100644
--- a/erts/emulator/beam/bif_instrs.tab
+++ b/erts/emulator/beam/bif_instrs.tab
@@ -432,9 +432,17 @@ nif_bif.call_nif() {
         live_hf_end = c_p->mbuf;
         ERTS_CHK_MBUF_SZ(c_p);
         erts_pre_nif(&env, c_p, (struct erl_module_nif*)I[2], NULL);
+
+        ASSERT((c_p->scheduler_data)->current_nif == NULL);
+        (c_p->scheduler_data)->current_nif = &env;
+
         nif_bif_result = (*fp)(&env, bif_nif_arity, reg);
         if (env.exception_thrown)
             nif_bif_result = THE_NON_VALUE;
+
+        ASSERT((c_p->scheduler_data)->current_nif == &env);
+        (c_p->scheduler_data)->current_nif = NULL;
+
         erts_post_nif(&env);
         ERTS_CHK_MBUF_SZ(c_p);
 
diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 0e0013e8a4..260656e2d3 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -388,12 +388,18 @@ erts_call_dirty_nif(ErtsSchedulerData *esdp, Process *c_p, BeamInstr *I, Eterm *
     erts_atomic32_read_band_mb(&c_p->state, ~(ERTS_PSFLG_DIRTY_CPU_PROC
 						   | ERTS_PSFLG_DIRTY_IO_PROC));
 
+    ASSERT(esdp->current_nif == NULL);
+    esdp->current_nif = &env;
+
     erts_proc_unlock(c_p, ERTS_PROC_LOCK_MAIN);
 
     result = (*dirty_nif)(&env, codemfa->arity, argv); /* Call dirty NIF */
 
     erts_proc_lock(c_p, ERTS_PROC_LOCK_MAIN);
 
+    ASSERT(esdp->current_nif == &env);
+    esdp->current_nif = NULL;
+
     ASSERT(env.proc->static_flags & ERTS_STC_FLG_SHADOW_PROC);
     ASSERT(env.proc->next == c_p);
 
@@ -4258,6 +4264,10 @@ int erts_nif_get_funcs(struct erl_module_nif* mod,
     return mod->entry.num_of_funcs;
 }
 
+Module *erts_nif_get_module(struct erl_module_nif *nif_mod) {
+    return nif_mod->mod;
+}
+
 Eterm erts_nif_call_function(Process *p, Process *tracee,
                              struct erl_module_nif* mod,
                              ErlNifFunc *fun, int argc, Eterm *argv)
diff --git a/erts/emulator/beam/erl_process.c b/erts/emulator/beam/erl_process.c
index 71541786d0..47b46fdd75 100644
--- a/erts/emulator/beam/erl_process.c
+++ b/erts/emulator/beam/erl_process.c
@@ -5698,6 +5698,7 @@ init_scheduler_data(ErtsSchedulerData* esdp, int num,
     esdp->ssi = ssi;
     esdp->current_process = NULL;
     esdp->current_port = NULL;
+    esdp->current_nif = NULL;
 
     esdp->virtual_reds = 0;
     esdp->cpu_id = -1;
diff --git a/erts/emulator/beam/erl_process.h b/erts/emulator/beam/erl_process.h
index e232776016..bb97b6690c 100644
--- a/erts/emulator/beam/erl_process.h
+++ b/erts/emulator/beam/erl_process.h
@@ -637,6 +637,7 @@ struct ErtsSchedulerData_ {
     ErtsSchedType type;
     Uint no;			/* Scheduler number for normal schedulers */
     Uint dirty_no;  /* Scheduler number for dirty schedulers */
+    struct enif_environment_t *current_nif;
     Process *dirty_shadow_process;
     Port *current_port;
     ErtsRunQueue *run_queue;
diff --git a/erts/emulator/beam/global.h b/erts/emulator/beam/global.h
index f4fea86cdc..2cf268162d 100644
--- a/erts/emulator/beam/global.h
+++ b/erts/emulator/beam/global.h
@@ -123,6 +123,7 @@ void erts_unload_nif(struct erl_module_nif* nif);
 extern void erl_nif_init(void);
 extern int erts_nif_get_funcs(struct erl_module_nif*,
                               struct enif_func_t **funcs);
+extern Module *erts_nif_get_module(struct erl_module_nif*);
 extern Eterm erts_nif_call_function(Process *p, Process *tracee,
                                     struct erl_module_nif*,
                                     struct enif_func_t *,