diff --git a/qemu/accel/tcg/cpu-exec.c b/qemu/accel/tcg/cpu-exec.c
index f938cfaa70..9c94efd111 100644
--- a/qemu/accel/tcg/cpu-exec.c
+++ b/qemu/accel/tcg/cpu-exec.c
@@ -288,7 +288,9 @@ static inline TranslationBlock *tb_find(CPUState *cpu,
     }
     /* See if we can patch the calling TB. */
     if (last_tb) {
+        tb_exec_unlock(cpu->uc);
         tb_add_jump(last_tb, tb_exit, tb);
+        tb_exec_lock(cpu->uc);
     }
 
     return tb;
diff --git a/qemu/include/tcg/tcg.h b/qemu/include/tcg/tcg.h
index 9908acaa57..fef0cc455a 100644
--- a/qemu/include/tcg/tcg.h
+++ b/qemu/include/tcg/tcg.h
@@ -709,6 +709,7 @@ struct TCGContext {
     TCGv_i64 cpu_bndu[4];
 
     /* qemu/tcg/i386/tcg-target.inc.c */
+    /* qemu/tcg/aarch64/tcg-target.inc.c */
     void *tb_ret_addr;
 
     /* target/riscv/translate.c */
diff --git a/qemu/softmmu/unicorn_vtlb.c b/qemu/softmmu/unicorn_vtlb.c
index 8b4e9e0b29..25c684064f 100644
--- a/qemu/softmmu/unicorn_vtlb.c
+++ b/qemu/softmmu/unicorn_vtlb.c
@@ -55,6 +55,8 @@ bool unicorn_fill_tlb(CPUState *cs, vaddr address, int size,
     struct hook *hook;
     HOOK_FOREACH_VAR_DECLARE;
 
+    cpu_restore_state(cs, retaddr, false);
+
     HOOK_FOREACH(uc, hook, UC_HOOK_TLB_FILL) {
         if (hook->to_delete) {
             continue;
diff --git a/qemu/tcg/aarch64/tcg-target.inc.c b/qemu/tcg/aarch64/tcg-target.inc.c
index c1f6e108b3..23349c7661 100644
--- a/qemu/tcg/aarch64/tcg-target.inc.c
+++ b/qemu/tcg/aarch64/tcg-target.inc.c
@@ -1858,8 +1858,6 @@ static void tcg_out_qemu_st(TCGContext *s, TCGReg data_reg, TCGReg addr_reg,
 #endif /* CONFIG_SOFTMMU */
 }
 
-static tcg_insn_unit *tb_ret_addr;
-
 static void tcg_out_op(TCGContext *s, TCGOpcode opc,
                        const TCGArg args[TCG_MAX_OP_ARGS],
                        const int const_args[TCG_MAX_OP_ARGS])
@@ -1885,7 +1883,7 @@ static void tcg_out_op(TCGContext *s, TCGOpcode opc,
             tcg_out_goto_long(s, s->code_gen_epilogue);
         } else {
             tcg_out_movi(s, TCG_TYPE_I64, TCG_REG_X0, a0);
-            tcg_out_goto_long(s, tb_ret_addr);
+            tcg_out_goto_long(s, s->tb_ret_addr);
         }
         break;
 
@@ -2859,7 +2857,7 @@ static void tcg_target_qemu_prologue(TCGContext *s)
     tcg_out_movi(s, TCG_TYPE_REG, TCG_REG_X0, 0);
 
     /* TB epilogue */
-    tb_ret_addr = s->code_ptr;
+    s->tb_ret_addr = s->code_ptr;
 
     /* Remove TCG locals stack space.  */
     tcg_out_insn(s, 3401, ADDI, TCG_TYPE_I64, TCG_REG_SP, TCG_REG_SP,
diff --git a/uc.c b/uc.c
index 0ff4196f1f..b44e5b7484 100644
--- a/uc.c
+++ b/uc.c
@@ -2428,6 +2428,7 @@ uc_err uc_context_restore(uc_engine *uc, uc_context *context)
         uc->snapshot_level = context->snapshot_level;
         ret = uc_restore_latest_snapshot(uc);
         if (ret != UC_ERR_OK) {
+            restore_jit_state(uc);
             return ret;
         }
         uc_snapshot(uc);
@@ -2441,9 +2442,12 @@ uc_err uc_context_restore(uc_engine *uc, uc_context *context)
     if (uc->context_content & UC_CTL_CONTEXT_CPU) {
         if (!uc->context_restore) {
             memcpy(uc->cpu->env_ptr, context->data, context->context_size);
+            restore_jit_state(uc);
             return UC_ERR_OK;
         } else {
-            return uc->context_restore(uc, context);
+            ret = uc->context_restore(uc, context);
+            restore_jit_state(uc);
+            return ret;
         }
     }
     return UC_ERR_OK;