8 files changed, 43 insertions, 16 deletions

diff --git a/src/lib_jit.c b/src/lib_jit.c
index a7c553ab4..af3e0a6ff 100644
--- a/src/lib_jit.c
+++ b/src/lib_jit.c
@@ -726,7 +726,7 @@ static void jit_init(lua_State *L)
 #if LJ_TARGET_UNALIGNED
   G(L)->tmptv.u64 = U64x(0000504d,4d500000);
 #endif
-  lj_dispatch_update(G(L));
+  lj_dispatch_update(G(L), 0);
 #if LJ_TARGET_UNALIGNED
   /* If you get a crash below then your toolchain indicates unaligned
   ** accesses are OK, but your kernel disagrees. I.e. fix your toolchain.
diff --git a/src/lj_dispatch.c b/src/lj_dispatch.c
index 636fb0338..9879a95b9 100644
--- a/src/lj_dispatch.c
+++ b/src/lj_dispatch.c
@@ -103,8 +103,11 @@ void lj_dispatch_init_hotcount(global_State *g)
 #define DISPMODE_PROF   0x40    /* Profiling active. */
 
 /* Update dispatch table depending on various flags. */
-void lj_dispatch_update(global_State *g)
+void LJ_FASTCALL lj_dispatch_update(global_State *g, int nolock)
 {
+#if LJ_HASPROFILE && !LJ_PROFILE_SIGPROF
+  int profile_locked = nolock ? 0 : lj_profile_lock();
+#endif
   uint8_t oldmode = g->dispatchmode;
   uint8_t mode = 0;
 #if LJ_HASJIT
@@ -208,6 +211,11 @@ void lj_dispatch_update(global_State *g)
       lj_dispatch_init_hotcount(g);
 #endif
   }
+#if LJ_HASPROFILE && !LJ_PROFILE_SIGPROF
+  if (profile_locked) lj_profile_unlock();
+#else
+  UNUSED(nolock);
+#endif
 }
 
 /* -- JIT mode setting ---------------------------------------------------- */
@@ -260,7 +268,7 @@ int luaJIT_setmode(lua_State *L, int idx, int mode)
         G2J(g)->flags &= ~(uint32_t)JIT_F_ON;
       else
         G2J(g)->flags |= (uint32_t)JIT_F_ON;
-      lj_dispatch_update(g);
+      lj_dispatch_update(g, 0);
     }
     break;
   case LUAJIT_MODE_FUNC:
@@ -335,7 +343,7 @@ LUA_API int lua_sethook(lua_State *L, lua_Hook func, int mask, int count)
   g->hookcount = g->hookcstart = (int32_t)count;
   g->hookmask = (uint8_t)((g->hookmask & ~HOOK_EVENTMASK) | mask);
   lj_trace_abort(g);  /* Abort recording on any hook change. */
-  lj_dispatch_update(g);
+  lj_dispatch_update(g, 0);
   return 1;
 }
 
diff --git a/src/lj_dispatch.h b/src/lj_dispatch.h
index 8492f7930..9baf762b2 100644
--- a/src/lj_dispatch.h
+++ b/src/lj_dispatch.h
@@ -132,7 +132,7 @@ LJ_FUNC void lj_dispatch_init(GG_State *GG);
 #if LJ_HASJIT
 LJ_FUNC void lj_dispatch_init_hotcount(global_State *g);
 #endif
-LJ_FUNC void lj_dispatch_update(global_State *g);
+LJ_FUNC void LJ_FASTCALL lj_dispatch_update(global_State *g, int nolock);
 
 /* Instruction dispatch callback for hooks or when recording. */
 LJ_FUNCA void LJ_FASTCALL lj_dispatch_ins(lua_State *L, const BCIns *pc);
diff --git a/src/lj_gc.c b/src/lj_gc.c
index 1fed0b54f..3ca6cf1bf 100644
--- a/src/lj_gc.c
+++ b/src/lj_gc.c
@@ -512,7 +512,7 @@ static void gc_call_finalizer(global_State *g, lua_State *L,
   TValue *top;
   lj_trace_abort(g);
   hook_entergc(g);  /* Disable hooks and new traces during __gc. */
-  if (LJ_HASPROFILE && (oldh & HOOK_PROFILE)) lj_dispatch_update(g);
+  if (LJ_HASPROFILE && (oldh & HOOK_PROFILE)) lj_dispatch_update(g, 0);
   g->gc.threshold = LJ_MAX_MEM;  /* Prevent GC steps. */
   top = VL->top;
   copyTV(VL, top++, mo);
@@ -522,7 +522,7 @@ static void gc_call_finalizer(global_State *g, lua_State *L,
   errcode = lj_vm_pcall(VL, top, 1+0, -1);  /* Stack: |mo|o| -> | */
   setgcref(g->cur_L, obj2gco(L));
   hook_restore(g, oldh);
-  if (LJ_HASPROFILE && (oldh & HOOK_PROFILE)) lj_dispatch_update(g);
+  if (LJ_HASPROFILE && (oldh & HOOK_PROFILE)) lj_dispatch_update(g, 0);
   g->gc.threshold = oldt;  /* Restore GC threshold. */
   if (errcode) {
     TValue tmp;
diff --git a/src/lj_profile.c b/src/lj_profile.c
index fcab1105b..5b2a2b6c7 100644
--- a/src/lj_profile.c
+++ b/src/lj_profile.c
@@ -118,6 +118,23 @@ void LJ_FASTCALL lj_profile_hook_leave(global_State *g)
     hook_leave(g);
   }
 }
+
+int lj_profile_lock(void)
+{
+  ProfileState *ps = &profile_state;
+  if (ps->g) {
+    profile_lock(ps);
+    return 1;
+  } else {
+    return 0;
+  }
+}
+
+void lj_profile_unlock(void)
+{
+  ProfileState *ps = &profile_state;
+  profile_unlock(ps);
+}
 #endif
 
 /* -- Profile callbacks --------------------------------------------------- */
@@ -134,14 +151,14 @@ void LJ_FASTCALL lj_profile_interpreter(lua_State *L)
     int samples = ps->samples;
     ps->samples = 0;
     g->hookmask = HOOK_VMEVENT;
-    lj_dispatch_update(g);
+    lj_dispatch_update(g, 1);
     profile_unlock(ps);
     ps->cb(ps->data, L, samples, ps->vmstate);  /* Invoke user callback. */
     profile_lock(ps);
     mask |= (g->hookmask & HOOK_PROFILE);
   }
   g->hookmask = mask;
-  lj_dispatch_update(g);
+  lj_dispatch_update(g, 1);
   profile_unlock(ps);
 }
 
@@ -160,7 +177,7 @@ static void profile_trigger(ProfileState *ps)
                   st == ~LJ_VMST_C ? 'C' :
                   st == ~LJ_VMST_GC ? 'G' : 'J';
     g->hookmask = (mask | HOOK_PROFILE);
-    lj_dispatch_update(g);
+    lj_dispatch_update(g, 1);
   }
   profile_unlock(ps);
 }
@@ -344,7 +361,7 @@ LUA_API void luaJIT_profile_stop(lua_State *L)
   if (G(L) == g) {  /* Only stop profiler if started by this VM. */
     profile_timer_stop(ps);
     g->hookmask &= ~HOOK_PROFILE;
-    lj_dispatch_update(g);
+    lj_dispatch_update(g, 0);
 #if LJ_HASJIT
     G2J(g)->prof_mode = 0;
     lj_trace_flushall(L);
diff --git a/src/lj_profile.h b/src/lj_profile.h
index 8b12437dd..43a36b6b1 100644
--- a/src/lj_profile.h
+++ b/src/lj_profile.h
@@ -14,6 +14,8 @@ LJ_FUNC void LJ_FASTCALL lj_profile_interpreter(lua_State *L);
 #if !LJ_PROFILE_SIGPROF
 LJ_FUNC void LJ_FASTCALL lj_profile_hook_enter(global_State *g);
 LJ_FUNC void LJ_FASTCALL lj_profile_hook_leave(global_State *g);
+LJ_FUNC int lj_profile_lock(void);
+LJ_FUNC void lj_profile_unlock(void);
 #endif
 
 #endif
diff --git a/src/lj_state.c b/src/lj_state.c
index 0c2c07505..0b997c6e8 100644
--- a/src/lj_state.c
+++ b/src/lj_state.c
@@ -337,7 +337,7 @@ LUA_API void lua_close(lua_State *L)
 #if LJ_HASJIT
   G2J(g)->flags &= ~JIT_F_ON;
   G2J(g)->state = LJ_TRACE_IDLE;
-  lj_dispatch_update(g);
+  lj_dispatch_update(g, 0);
 #endif
   for (i = 0;;) {
     hook_enter(g);
diff --git a/src/lj_trace.c b/src/lj_trace.c
index e55045f4e..24a1f7e45 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -656,7 +656,7 @@ static int trace_abort(jit_State *J)
   } else if (e == LJ_TRERR_MCODEAL) {
     if (!J->mcarea) {  /* Disable JIT compiler if first mcode alloc fails. */
       J->flags &= ~JIT_F_ON;
-      lj_dispatch_update(J2G(J));
+      lj_dispatch_update(J2G(J), 0);
     }
     lj_trace_flushall(L);
   }
@@ -687,7 +687,7 @@ static TValue *trace_state(lua_State *L, lua_CFunction dummy, void *ud)
     case LJ_TRACE_START:
       J->state = LJ_TRACE_RECORD;  /* trace_start() may change state. */
       trace_start(J);
-      lj_dispatch_update(J2G(J));
+      lj_dispatch_update(J2G(J), 0);
       if (J->state != LJ_TRACE_RECORD_1ST)
         break;
       /* fallthrough */
@@ -745,7 +745,7 @@ static TValue *trace_state(lua_State *L, lua_CFunction dummy, void *ud)
       trace_stop(J);
       setvmstate(J2G(J), INTERP);
       J->state = LJ_TRACE_IDLE;
-      lj_dispatch_update(J2G(J));
+      lj_dispatch_update(J2G(J), 0);
       return NULL;
 
     default:  /* Trace aborted asynchronously. */
@@ -757,7 +757,7 @@ static TValue *trace_state(lua_State *L, lua_CFunction dummy, void *ud)
         goto retry;
       setvmstate(J2G(J), INTERP);
       J->state = LJ_TRACE_IDLE;
-      lj_dispatch_update(J2G(J));
+      lj_dispatch_update(J2G(J), 0);
       return NULL;
     }
   } while (J->state > LJ_TRACE_RECORD);