[Patch v4 17/46] UefiCpuPkg/MpInitLib: Register one End of PEI callback function

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

In PeiMpInitLib, register End of PEI callback function CpuMpEndOfPeiCallback().

Cc: Michael Kinney <michael.d.kinney@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Giri P Mudusuru <giri.p.mudusuru@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.h    | 20 +++++++++++++++++
 UefiCpuPkg/Library/MpInitLib/PeiMpLib.c | 40 +++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 5525a96..a2818e6 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -265,5 +265,25 @@ MicrocodeDetect (
   IN CPU_MP_DATA             *CpuMpData
   );
 
+/**
+  Notify function on End Of PEI PPI.
+
+  On S3 boot, this function will restore wakeup buffer data.
+  On normal boot, this function will flag wakeup buffer to be un-used type.
+
+  @param[in]  PeiServices        The pointer to the PEI Services Table.
+  @param[in]  NotifyDescriptor   Address of the notification descriptor data structure.
+  @param[in]  Ppi                Address of the PPI that was installed.
+
+  @retval EFI_SUCCESS        When everything is OK.
+**/
+EFI_STATUS
+EFIAPI
+CpuMpEndOfPeiCallback (
+  IN EFI_PEI_SERVICES             **PeiServices,
+  IN EFI_PEI_NOTIFY_DESCRIPTOR    *NotifyDescriptor,
+  IN VOID                         *Ppi
+  );
+
 #endif
 
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
index 6211e71..f37e985 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpLib.c
@@ -13,6 +13,14 @@
 **/
 
 #include "MpLib.h"
+#include <Ppi/EndOfPeiPhase.h>
+#include <Library/PeiServicesLib.h>
+
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_PEI_NOTIFY_DESCRIPTOR mMpInitLibNotifyList = {
+  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
+  &gEfiEndOfPeiSignalPpiGuid,
+  CpuMpEndOfPeiCallback
+};
 
 /**
   Get pointer to CPU MP Data structure.
@@ -55,6 +63,32 @@ SaveCpuMpData (
 
 
 /**
+/**
+  Notify function on End Of PEI PPI.
+
+  On S3 boot, this function will restore wakeup buffer data.
+  On normal boot, this function will flag wakeup buffer to be un-used type.
+
+  @param[in]  PeiServices        The pointer to the PEI Services Table.
+  @param[in]  NotifyDescriptor   Address of the notification descriptor data structure.
+  @param[in]  Ppi                Address of the PPI that was installed.
+
+  @retval EFI_SUCCESS        When everything is OK.
+**/
+EFI_STATUS
+EFIAPI
+CpuMpEndOfPeiCallback (
+  IN EFI_PEI_SERVICES             **PeiServices,
+  IN EFI_PEI_NOTIFY_DESCRIPTOR    *NotifyDescriptor,
+  IN VOID                         *Ppi
+  )
+{
+
+  DEBUG ((DEBUG_INFO, "PeiMpInitLib: CpuMpEndOfPeiCallback () invoked\n"));
+
+  return EFI_SUCCESS;
+}
+/**
   Initialize global data for MP support.
 
   @param[in] CpuMpData  The pointer to CPU MP Data structure.
@@ -64,8 +98,14 @@ InitMpGlobalData (
   IN CPU_MP_DATA               *CpuMpData
   )
 {
+  EFI_STATUS      Status;
 
   SaveCpuMpData (CpuMpData);
+  //
+  // Register an event for EndOfPei
+  //
+  Status  = PeiServicesNotifyPpi (&mMpInitLibNotifyList);
+  ASSERT_EFI_ERROR (Status);
 }
 
 /**
-- 
2.7.4.windows.1