[Patch] UefiCpuPkg/MpInitLib: Rollback old change 2a5997f8.

[Patch] UefiCpuPkg/MpInitLib: Rollback old change 2a5997f8.

[Eric Dong]

In some special cases, after BSP send Init-sipi-sipi signal
AP need more time to start the Ap procedure. In this case
BSP may think AP has finish its task but truly AP not begin
yet.

Rollback former change to keep the status which only used
when AP truly finished task.

Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.c | 10 ++++++----
 UefiCpuPkg/Library/MpInitLib/MpLib.h |  1 +
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 85ca4a2946..7f4d6e60bd 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -696,7 +696,7 @@ ApWakeupFunction (
             }
           }
         }
-        SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateIdle);
+        SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateFinished);
       }
     }
 
@@ -1370,10 +1370,11 @@ CheckThisAP (
   //
   // If the AP finishes for StartupThisAP(), return EFI_SUCCESS.
   //
-  if (GetApState(CpuData) == CpuStateIdle) {
+  if (GetApState(CpuData) == CpuStateFinished) {
     if (CpuData->Finished != NULL) {
       *(CpuData->Finished) = TRUE;
     }
+    SetApState (CpuData, CpuStateIdle);
     return EFI_SUCCESS;
   } else {
     //
@@ -1434,9 +1435,10 @@ CheckAllAPs (
     // Only BSP and corresponding AP access this unit of CPU Data. This means the AP will not modify the
     // value of state after setting the it to CpuStateIdle, so BSP can safely make use of its value.
     //
-    if (GetApState(CpuData) == CpuStateIdle) {
+    if (GetApState(CpuData) == CpuStateFinished) {
       CpuMpData->RunningCount --;
       CpuMpData->CpuData[ProcessorNumber].Waiting = FALSE;
+      SetApState(CpuData, CpuStateIdle);
 
       //
       // If in Single Thread mode, then search for the next waiting AP for execution.
@@ -1937,7 +1939,7 @@ SwitchBSPWorker (
   //
   // Wait for old BSP finished AP task
   //
-  while (GetApState (&CpuMpData->CpuData[CallerNumber]) != CpuStateIdle) {
+  while (GetApState (&CpuMpData->CpuData[CallerNumber]) != CpuStateFinished) {
     CpuPause ();
   }
 
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 773db76b61..5f6986c240 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -89,6 +89,7 @@ typedef enum {
   CpuStateIdle,
   CpuStateReady,
   CpuStateBusy,
+  CpuStateFinished,
   CpuStateDisabled
 } CPU_STATE;
 
-- 
2.15.0.windows.1