[PATCH v2 1/1] UefiCpuPkg: Extend SMM CPU Service with rendezvous support.

["Li, Zhihao" <zhihao.li@intel.com>]

REF： https://bugzilla.tianocore.org/show_bug.cgi?id=3815

This patch define a new Protocol with the new services
SmmWaitForAllProcessor(), which can be used by SMI handler
to optionally wait for other APs to complete SMM rendezvous in
relaxed AP mode.

A new library SmmCpuRendezvousLib is provided to abstract the service
into library API to simple SMI handler code.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Siyuan Fu <siyuan.fu@intel.com>

Signed-off-by: Zhihao Li <zhihao.li@intel.com>
---
 UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c   | 95 ++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c                         | 66 +++++++++++++-
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c                          | 14 ++-
 UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h               | 27 ++++++
 UefiCpuPkg/Include/Protocol/SmmCpuService.h                    | 36 +++++++-
 UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf | 32 +++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h                     | 30 ++++++-
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf                   |  5 +-
 UefiCpuPkg/UefiCpuPkg.dec                                      |  5 +-
 9 files changed, 300 insertions(+), 10 deletions(-)

diff --git a/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c
new file mode 100644
index 000000000000..03e507bf6b52
--- /dev/null
+++ b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c
@@ -0,0 +1,95 @@
+/** @file
+SMM CPU Rendezvous library header file.
+
+Copyright (c) 2021 - 2022, Intel Corporation. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+
+#include <Base.h>
+#include <Uefi.h>
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MmServicesTableLib.h>
+#include <Protocol/SmmCpuService.h>
+#include <Library/SmmCpuRendezvousLib.h>
+
+STATIC EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL  *mSmmCpuRendezvous = NULL;
+STATIC VOID                               *mSmmCpuRendezvousRegistration = NULL;
+
+/**
+  Register status code callback function only when Report Status Code protocol
+  is installed.
+
+  @param Protocol       Points to the protocol's unique identifier.
+  @param Interface      Points to the interface instance.
+  @param Handle         The handle on which the interface was installed.
+
+  @retval EFI_SUCCESS   Notification runs successfully.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmCpuServiceProtocolNotify (
+  IN CONST EFI_GUID        *Protocol,
+  IN VOID                  *Interface,
+  IN EFI_HANDLE            Handle
+  )
+{
+  EFI_STATUS                   Status;
+
+  Status = gMmst->MmLocateProtocol (
+                    &gEdkiiSmmCpuRendezvousProtocolGuid,
+                    NULL,
+                    (VOID **) &mSmmCpuRendezvous
+                    );
+  ASSERT_EFI_ERROR (Status);
+
+  return EFI_SUCCESS;
+}
+
+/**
+  This routine wait for all AP processors to arrive in SMM.
+
+  @param  BlockingMode  Blocking mode or non-blocking mode.
+
+  @retval EFI_SUCCESS            All avaiable APs arrived.
+  @retval EFI_TIMEOUT            Wait for all APs until timeout.
+  @retval Other                  Fail to register Smm cpu rendezvous services notify.
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  BOOLEAN  BlockingMode
+  )
+{
+  EFI_STATUS  Status;
+
+  if (mSmmCpuRendezvousRegistration == NULL && mSmmCpuRendezvous == NULL) {
+    //
+    // locate Smm cpu rendezvous protocol for the first time execute the function.
+    //
+    Status = gMmst->MmLocateProtocol (&gEdkiiSmmCpuRendezvousProtocolGuid, NULL, (VOID **) &mSmmCpuRendezvous);
+    if (EFI_ERROR (Status)) {
+      Status = gMmst->MmRegisterProtocolNotify (
+                &gEdkiiSmmCpuRendezvousProtocolGuid,
+                SmmCpuServiceProtocolNotify,
+                &mSmmCpuRendezvousRegistration
+                );
+      if (EFI_ERROR (Status)) {
+        return Status;
+      }
+    }
+  }
+
+  if (mSmmCpuRendezvous == NULL) {
+    return EFI_SUCCESS;
+  }
+
+  Status = mSmmCpuRendezvous->WaitForAllProcessor (
+              mSmmCpuRendezvous,
+              BlockingMode
+              );
+  return Status;
+}
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
index 5d624f8e9ed6..85c3c5c15a26 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
@@ -1,7 +1,7 @@
 /** @file
 Implementation of SMM CPU Services Protocol.
 
-Copyright (c) 2011 - 2015, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2011 - 2022, Intel Corporation. All rights reserved.<BR>
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -20,6 +20,13 @@ EFI_SMM_CPU_SERVICE_PROTOCOL  mSmmCpuService = {
   SmmRegisterExceptionHandler
 };
 
+//
+// EDKII SMM CPU Rendezvous Service Protocol instance
+//
+EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL  mEdkiiSmmCpuRendezvousService = {
+  SmmWaitForAllProcessor
+};
+
 /**
   Gets processor information on the requested processor at the instant this call is made.
 
@@ -350,6 +357,7 @@ SmmRegisterExceptionHandler (
   @param ImageHandle The firmware allocated handle for the EFI image.
 
   @retval EFI_SUCCESS    EFI SMM CPU Services Protocol was installed successfully.
+  @retval Other          Fail to install protocol.
 **/
 EFI_STATUS
 InitializeSmmCpuServices (
@@ -364,6 +372,62 @@ InitializeSmmCpuServices (
                     EFI_NATIVE_INTERFACE,
                     &mSmmCpuService
                     );
+  if (EFI_ERROR (Status)){
+    return Status;
+  }
+
+  Status = gSmst->SmmInstallProtocolInterface (
+                    &Handle,
+                    &gEdkiiSmmCpuRendezvousProtocolGuid,
+                    EFI_NATIVE_INTERFACE,
+                    &mEdkiiSmmCpuRendezvousService
+                    );
   ASSERT_EFI_ERROR (Status);
   return Status;
 }
+
+/**
+  Wait for all processors enterring SMM until all CPUs are already synchronized or not.
+
+  If BlockingMode is False, timeout value is zero.
+
+  @param  This                   A pointer to the EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL instance.
+  @param  BlockingMode           Blocking mode or non-blocking mode.
+
+  @retval EFI_SUCCESS    All avaiable APs arrived.
+  @retval EFI_TIMEOUT    Wait for all APs until timeout.
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL  *This,
+  IN  BOOLEAN                           BlockingMode
+  )
+{
+  EFI_STATUS                Status;
+
+  //
+  // Return success immediately if all CPUs are already synchronized.
+  //
+  if (mSmmMpSyncData->AllApArrivedWithException) {
+    Status = EFI_SUCCESS;
+    goto ON_EXIT;
+  }
+
+  if (!BlockingMode) {
+    Status = EFI_TIMEOUT;
+    goto ON_EXIT;
+  }
+
+  //
+  // There are some APs outside SMM, Wait for all avaiable APs to arrive.
+  //
+  SmmWaitForApArrival (BlockingMode);
+  Status = mSmmMpSyncData->AllApArrivedWithException ? EFI_SUCCESS : EFI_TIMEOUT;
+
+ON_EXIT:
+  if (!mSmmMpSyncData->AllApArrivedWithException){
+    DEBUG ((DEBUG_INFO, "EdkiiSmmWaitForAllApArrival: Timeout to wait all APs arrival\n"));
+  }
+  return Status;
+}
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 882dee4fe246..4051a9dba179 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -1,7 +1,7 @@
 /** @file
 SMM MP service implementation
 
-Copyright (c) 2009 - 2021, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2009 - 2022, Intel Corporation. All rights reserved.<BR>
 Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 
 SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -261,7 +261,7 @@ IsLmceSignaled (
 **/
 VOID
 SmmWaitForApArrival (
-  VOID
+  IN  BOOLEAN  BlockingMode
   )
 {
   UINT64   Timer;
@@ -271,6 +271,13 @@ SmmWaitForApArrival (
 
   ASSERT (*mSmmMpSyncData->Counter <= mNumberOfCpus);
 
+  //
+  // if block is False, do not wait and return immediately.
+  //
+  if (!BlockingMode){
+    return;
+  }
+
   LmceEn     = FALSE;
   LmceSignal = FALSE;
   if (mMachineCheckSupported) {
@@ -511,7 +518,7 @@ BSPHandler (
     //
     // Wait for APs to arrive
     //
-    SmmWaitForApArrival ();
+    SmmWaitForApArrival (TRUE);
 
     //
     // Lock the counter down and retrieve the number of APs
@@ -1886,6 +1893,7 @@ InitializeMpSyncData (
     *mSmmMpSyncData->Counter       = 0;
     *mSmmMpSyncData->InsideSmm     = FALSE;
     *mSmmMpSyncData->AllCpusInSync = FALSE;
+    mSmmMpSyncData->AllApArrivedWithException = FALSE;
 
     for (CpuIndex = 0; CpuIndex < gSmmCpuPrivate->SmmCoreEntryContext.NumberOfCpus; CpuIndex++) {
       mSmmMpSyncData->CpuData[CpuIndex].Busy =
diff --git a/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h b/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h
new file mode 100644
index 000000000000..6f90b3146445
--- /dev/null
+++ b/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h
@@ -0,0 +1,27 @@
+/** @file
+SMM CPU Rendezvous library header file.
+
+Copyright (c) 2021 - 2022, Intel Corporation. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#ifndef _SMM_CPU_RENDEZVOUS_H_
+#define _SMM_CPU_RENDEZVOUS_H_
+
+/**
+  This routine wait for all AP processors to arrive in SMM.
+
+  @param  BlockingMode  Blocking mode or non-blocking mode.
+
+  @retval EFI_SUCCESS   All processors checked in to SMM.
+  @retval EFI_TIMEOUT   Wait for all APs until timeout.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  BOOLEAN BlockingMode
+  );
+
+#endif
diff --git a/UefiCpuPkg/Include/Protocol/SmmCpuService.h b/UefiCpuPkg/Include/Protocol/SmmCpuService.h
index 952767afce75..95376e87e933 100644
--- a/UefiCpuPkg/Include/Protocol/SmmCpuService.h
+++ b/UefiCpuPkg/Include/Protocol/SmmCpuService.h
@@ -1,7 +1,7 @@
 /** @file
 SMM CPU Service protocol definition.
 
-Copyright (c) 2013 - 2015, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2013 - 2022, Intel Corporation. All rights reserved.<BR>
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -200,4 +200,38 @@ struct _EFI_SMM_CPU_SERVICE_PROTOCOL {
 
 extern EFI_GUID  gEfiSmmCpuServiceProtocolGuid;
 
+//
+//  EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL provide SMM CPU
+//  rendezvous service support.
+//
+#define EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL_GUID \
+  { \
+    0xaa00d50b, 0x4911, 0x428f, {0xb9, 0x1a, 0xa5, 0x9d, 0xdb, 0x13, 0xe2, 0x4c} \
+  }
+
+typedef struct _EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL;
+
+/**
+  Wait for all APs to arrive SMM mode in given timeout constraint.
+
+  @param  This                  A pointer to the EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL instance.
+  @param  BlockingMode          Block or non-block mode.
+
+  @retval EFI_SUCCESS           All APs have arrived SMM mode except SMI disabled APs.
+  @retval EFI_TIMEOUT           There are APs not in SMM mode in given timeout constraint.
+
+**/
+typedef
+EFI_STATUS
+(EFIAPI *EDKII_WAIT_FOR_ALL_PROCESSOR) (
+  IN  EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL  *This,
+  IN  BOOLEAN                           BlockingMode
+  );
+
+struct _EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL {
+  EDKII_WAIT_FOR_ALL_PROCESSOR        WaitForAllProcessor;
+};
+
+extern EFI_GUID gEdkiiSmmCpuRendezvousProtocolGuid;
+
 #endif
diff --git a/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf
new file mode 100644
index 000000000000..f4b9afb05f15
--- /dev/null
+++ b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf
@@ -0,0 +1,32 @@
+## @file
+# Component description file for CPU SMM Rendezvous check library
+#
+# Copyright (c) 2021 - 2022, Intel Corporation. All rights reserved.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = SmmCpuRendezvousLib
+  FILE_GUID                      = 1509Bb36-9Ba4-438B-B195-Ac5914Db14E2
+  MODULE_TYPE                    = DXE_SMM_DRIVER
+  LIBRARY_CLASS                  = SmmCpuRendezvousLib|MM_STANDALONE DXE_SMM_DRIVER
+
+[Sources]
+  SmmCpuRendezvousLib.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  UefiCpuPkg/UefiCpuPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+  MmServicesTableLib
+
+[Pcd]
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmApSyncTimeout                 ## CONSUMES
+
+[Protocols]
+  gEdkiiSmmCpuRendezvousProtocolGuid
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 26d07c5b5ea0..96f98cf85c25 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -1,7 +1,7 @@
 /** @file
 Agent Module to load other modules to deploy SMM Entry Vector for X86 CPU.
 
-Copyright (c) 2009 - 2020, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2009 - 2022, Intel Corporation. All rights reserved.<BR>
 Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 
 SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -428,6 +428,7 @@ typedef struct {
   volatile SMM_CPU_SYNC_MODE    EffectiveSyncMode;
   volatile BOOLEAN              SwitchBsp;
   volatile BOOLEAN              *CandidateBsp;
+  volatile BOOLEAN              AllApArrivedWithException;
   EFI_AP_PROCEDURE              StartupProcedure;
   VOID                          *StartupProcArgs;
 } SMM_DISPATCHER_MP_SYNC_DATA;
@@ -1488,4 +1489,31 @@ IsRestrictedMemoryAccess (
   VOID
   );
 
+/**
+  Choose blocking or non-blocking mode to Wait for all APs
+
+  @param  This                  A pointer to the EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL instance.
+  @param  BlockingMode          Blocking or non-blocking mode
+
+  @retval EFI_SUCCESS           All APs have arrived SMM mode except SMI disabled APs.
+  @retval EFI_TIMEOUT           There are APs not in SMM mode in given timeout constraint.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  EDKII_SMM_CPU_RENDEZVOUS_PROTOCOL  *This,
+  IN  BOOLEAN                           BlockingMode
+  );
+
+/**
+  Choose blocking or non-blocking mode to wait for all APs. True for Blocking and false for not.
+  Insure when this function returns, no AP will execute normal mode code before entering SMM, except SMI disabled APs.
+
+**/
+VOID
+SmmWaitForApArrival (
+  IN  BOOLEAN  BlockingMode
+  );
+
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
index 0e88071c7079..deef00f9c6e9 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
@@ -4,7 +4,7 @@
 # This SMM driver performs SMM initialization, deploy SMM Entry Vector,
 # provides CPU specific services in SMM.
 #
-# Copyright (c) 2009 - 2021, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2009 - 2022, Intel Corporation. All rights reserved.<BR>
 # Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 #
 # SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -107,7 +107,8 @@
   gEfiSmmReadyToLockProtocolGuid           ## NOTIFY
   gEfiSmmCpuServiceProtocolGuid            ## PRODUCES
   gEdkiiSmmMemoryAttributeProtocolGuid     ## PRODUCES
-  gEfiMmMpProtocolGuid                    ## PRODUCES
+  gEfiMmMpProtocolGuid                     ## PRODUCES
+  gEdkiiSmmCpuRendezvousProtocolGuid       ## PRODUCES
 
 [Guids]
   gEfiAcpiVariableGuid                     ## SOMETIMES_CONSUMES ## HOB # it is used for S3 boot.
diff --git a/UefiCpuPkg/UefiCpuPkg.dec b/UefiCpuPkg/UefiCpuPkg.dec
index 7de66fde674c..1951eb294c6c 100644
--- a/UefiCpuPkg/UefiCpuPkg.dec
+++ b/UefiCpuPkg/UefiCpuPkg.dec
@@ -1,7 +1,7 @@
 ## @file  UefiCpuPkg.dec
 # This Package provides UEFI compatible CPU modules and libraries.
 #
-# Copyright (c) 2007 - 2021, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2007 - 2022, Intel Corporation. All rights reserved.<BR>
 #
 # SPDX-License-Identifier: BSD-2-Clause-Patent
 #
@@ -77,7 +77,8 @@
 
 [Protocols]
   ## Include/Protocol/SmmCpuService.h
-  gEfiSmmCpuServiceProtocolGuid  = { 0x1d202cab, 0xc8ab, 0x4d5c, { 0x94, 0xf7, 0x3c, 0xfc, 0xc0, 0xd3, 0xd3, 0x35 }}
+  gEfiSmmCpuServiceProtocolGuid   = { 0x1d202cab, 0xc8ab, 0x4d5c, { 0x94, 0xf7, 0x3c, 0xfc, 0xc0, 0xd3, 0xd3, 0x35 }}
+  gEdkiiSmmCpuRendezvousProtocolGuid = { 0xaa00d50b, 0x4911, 0x428f, { 0xb9, 0x1a, 0xa5, 0x9d, 0xdb, 0x13, 0xe2, 0x4c }}
 
   ## Include/Protocol/SmMonitorInit.h
   gEfiSmMonitorInitProtocolGuid  = { 0x228f344d, 0xb3de, 0x43bb, { 0xa4, 0xd7, 0xea, 0x20, 0xb, 0x1b, 0x14, 0x82 }}
-- 
2.26.2.windows.1