[Patch 2/6] UefiCpuPkg/MpInitLib: Move allocating reserved memory for AP loop code

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

In Exit Boot Services callback function, we cannot use allocate memory services
because it may change the memory map that has been gotten by OS.

This fix is to move allocating reserved memory for AP loop code to
InitMpGlobalData() and save the memory address in one global variable.

Cc: Michael Kinney <michael.d.kinney@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 42d320f..383eec9 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -23,7 +23,7 @@ CPU_MP_DATA      *mCpuMpData = NULL;
 EFI_EVENT        mCheckAllApsEvent = NULL;
 EFI_EVENT        mMpInitExitBootServicesEvent = NULL;
 volatile BOOLEAN mStopCheckAllApsStatus = TRUE;
-
+VOID             *mReservedApLoopFunc = NULL;
 
 /**
   Get the pointer to CPU MP Data structure.
@@ -258,19 +258,11 @@ MpInitExitBootServicesCallback (
   )
 {
   CPU_MP_DATA               *CpuMpData;
-  VOID                      *ReservedApLoopFunc;
-  //
-  // Avoid APs access invalid buff data which allocated by BootServices,
-  // so we will allocate reserved data for AP loop code.
-  //
+
   CpuMpData = GetCpuMpData ();
   CpuMpData->PmCodeSegment = GetProtectedModeCS ();
   CpuMpData->ApLoopMode = PcdGet8 (PcdCpuApLoopMode);
-  ReservedApLoopFunc = AllocateReservedCopyPool (
-                         CpuMpData->AddressMap.RelocateApLoopFuncSize,
-                         CpuMpData->AddressMap.RelocateApLoopFuncAddress
-                         );
-  WakeUpAP (CpuMpData, TRUE, 0, RelocateApLoop, ReservedApLoopFunc);
+  WakeUpAP (CpuMpData, TRUE, 0, RelocateApLoop, mReservedApLoopFunc);
   DEBUG ((DEBUG_INFO, "MpInitExitBootServicesCallback() done!\n"));
 }
 
@@ -288,6 +280,18 @@ InitMpGlobalData (
 
   SaveCpuMpData (CpuMpData);
 
+  //
+  // Avoid APs access invalid buff data which allocated by BootServices,
+  // so we will allocate reserved data for AP loop code.
+  // Allocating it in advance since memory services are not available in
+  // Exit Boot Services callback function.
+  //
+  mReservedApLoopFunc = AllocateReservedCopyPool (
+                          CpuMpData->AddressMap.RelocateApLoopFuncSize,
+                          CpuMpData->AddressMap.RelocateApLoopFuncAddress
+                          );
+  ASSERT (mReservedApLoopFunc != NULL);
+
   Status = gBS->CreateEvent (
                   EVT_TIMER | EVT_NOTIFY_SIGNAL,
                   TPL_NOTIFY,
-- 
2.7.4.windows.1