[PATCH v2 3/3] UefiCpuPkg/PiSmmCpuDxeSmm: Decrease mNumberToFinish in AP safe code

[Jeff Fan <jeff.fan@intel.com>]

We will put APs into hlt-loop in safe code. But we decrease mNumberToFinish
before APs enter into the safe code. Paolo pointed out this gap.

This patch is to move mNumberToFinish decreasing to the safe code. It could
make sure BSP could wait for all APs are running in safe code.

https://bugzilla.tianocore.org/show_bug.cgi?id=216

Reported-by: Paolo Bonzini <pbonzini@redhat.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Paolo Bonzini <pbonzini@redhat.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c             | 17 +++++++----------
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmFuncsArch.c |  6 ++++--
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h    |  4 +++-
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c  |  6 ++++--
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
index e53e096..34547e0 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuS3.c
@@ -79,9 +79,11 @@ BOOLEAN                      mAcpiS3Enable = TRUE;
 
 UINT8                        *mApHltLoopCode = NULL;
 UINT8                        mApHltLoopCodeTemplate[] = {
-                               0xFA,        // cli
-                               0xF4,        // hlt
-                               0xEB, 0xFC   // jmp $-2
+                               0x8B, 0x44, 0x24, 0x04,  // mov  eax, dword ptr [esp+4]
+                               0xF0, 0xFF, 0x08,        // lock dec  dword ptr [eax]
+                               0xFA,                    // cli
+                               0xF4,                    // hlt
+                               0xEB, 0xFC               // jmp $-2
                                };
 
 /**
@@ -399,17 +401,12 @@ MPRendezvousProcedure (
   }
 
   //
-  // Count down the number with lock mechanism.
-  //
-  InterlockedDecrement (&mNumberToFinish);
-
-  //
-  // Place AP into the safe code
+  // Place AP into the safe code, count down the number with lock mechanism in the safe code.
   //
   TopOfStack  = (UINT32) (UINTN) Stack + sizeof (Stack);
   TopOfStack &= ~(UINT32) (CPU_STACK_ALIGNMENT - 1);
   CopyMem ((VOID *) (UINTN) mApHltLoopCode, mApHltLoopCodeTemplate, sizeof (mApHltLoopCodeTemplate));
-  TransferApToSafeState ((UINT32) (UINTN) mApHltLoopCode, TopOfStack);
+  TransferApToSafeState ((UINT32) (UINTN) mApHltLoopCode, TopOfStack, &mNumberToFinish);
 }
 
 /**
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmFuncsArch.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmFuncsArch.c
index 8b880d6..9760373 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmFuncsArch.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmFuncsArch.c
@@ -100,17 +100,19 @@ InitGdt (
 
   @param[in] ApHltLoopCode    The 32-bit address of the safe hlt-loop function.
   @param[in] TopOfStack       A pointer to the new stack to use for the ApHltLoopCode.
+  @param[in] NumberToFinish   Semaphore of APs finish count.
 
 **/
 VOID
 TransferApToSafeState (
   IN UINT32             ApHltLoopCode,
-  IN UINT32             TopOfStack
+  IN UINT32             TopOfStack,
+  IN UINT32             *NumberToFinish
   )
 {
   SwitchStack (
     (SWITCH_STACK_ENTRY_POINT) (UINTN) ApHltLoopCode,
-    NULL,
+    NumberToFinish,
     NULL,
     (VOID *) (UINTN) TopOfStack
     );
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 6c98659..88d9c85 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -830,12 +830,14 @@ GetAcpiS3EnableFlag (
 
   @param[in] ApHltLoopCode    The 32-bit address of the safe hlt-loop function.
   @param[in] TopOfStack       A pointer to the new stack to use for the ApHltLoopCode.
+  @param[in] NumberToFinish   Semaphore of APs finish count.
 
 **/
 VOID
 TransferApToSafeState (
   IN UINT32             ApHltLoopCode,
-  IN UINT32             TopOfStack
+  IN UINT32             TopOfStack,
+  IN UINT32             *NumberToFinish
   );
 
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c
index 62338b7..6844c3f 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmFuncsArch.c
@@ -105,18 +105,20 @@ GetProtectedModeCS (
 
   @param[in] ApHltLoopCode    The 32-bit address of the safe hlt-loop function.
   @param[in] TopOfStack       A pointer to the new stack to use for the ApHltLoopCode.
+  @param[in] NumberToFinish   Semaphore of APs finish count.
 
 **/
 VOID
 TransferApToSafeState (
   IN UINT32             ApHltLoopCode,
-  IN UINT32             TopOfStack
+  IN UINT32             TopOfStack,
+  IN UINT32             *NumberToFinish
   )
 {
   AsmDisablePaging64 (
     GetProtectedModeCS (),
     (UINT32) (UINTN) ApHltLoopCode,
-    0,
+    (UINT32) (UINTN) NumberToFinish,
     0,
     TopOfStack
     );
-- 
2.9.3.windows.2