1 files changed, 7 insertions, 2 deletions
diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/UefiCpuPkg/CpuDxe/CpuMp.c
index 25b6fcb056..25c909193e 100644
--- a/UefiCpuPkg/CpuDxe/CpuMp.c
+++ b/UefiCpuPkg/CpuDxe/CpuMp.c
@@ -1176,11 +1176,14 @@ ProcessorToIdleState (
   }
 
   //
-  // Avoid forcibly reset AP caused the AP State is not updated.
+  // Avoid forcibly reset AP caused the timeout AP State is not
+  // updated.
   //
   GetMpSpinLock (CpuData);
+  if (CpuData->State == CpuStateBusy) {
+    CpuData->Procedure = NULL;
+  }
   CpuData->State = CpuStateIdle;
-  CpuData->Procedure = NULL;
   ReleaseMpSpinLock (CpuData);
 
   while (TRUE) {
@@ -1190,6 +1193,8 @@ ProcessorToIdleState (
     ReleaseMpSpinLock (CpuData);
 
     if (Procedure != NULL) {
+      SetApState (CpuData, CpuStateBusy);
+
       Procedure (ProcedureArgument);
 
       GetMpSpinLock (CpuData);