[PATCH V2 5/6] MdeModulePkg/SmmPerformanceLib: Disable perf-logging after ExitBS

["Ni, Ray" <ray.ni@intel.com>]

Because SMM perf-logging is migrated to non-SMRAM at ReadyToBoot
by DxeCorePerformanceLib, the perf-logging after ExitBS is useless and
impact the SMI latency at runtime.
Hence the SmmPerformanceLib is updated to disable perf-logging
after ExitBS.

Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Dandan Bi <dandan.bi@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Jian J Wang <jian.j.wang@intel.com>
---
 .../SmmPerformanceLib/SmmPerformanceLib.c     | 63 ++++++++++++++++++-
 .../SmmPerformanceLib/SmmPerformanceLib.inf   |  4 ++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.c b/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.c
index 623f8a978c..b9c33c0f64 100644
--- a/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.c
+++ b/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.c
@@ -23,6 +23,36 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 // The cached SMM Performance Protocol and SMM PerformanceEx Protocol interface.
 EDKII_PERFORMANCE_MEASUREMENT_PROTOCOL  *mPerformanceMeasurement = NULL;
 BOOLEAN                                 mPerformanceMeasurementEnabled;
+VOID                                    *mPerformanceLibExitBootServicesRegistration;
+
+/**
+  This is the Event call back function is triggered in SMM to notify the Library
+  the system is entering runtime phase.
+
+  @param[in] Protocol   Points to the protocol's unique identifier
+  @param[in] Interface  Points to the interface instance
+  @param[in] Handle     The handle on which the interface was installed
+
+  @retval EFI_SUCCESS SmmAtRuntimeCallBack runs successfully
+ **/
+EFI_STATUS
+EFIAPI
+SmmPerformanceLibExitBootServicesCallback (
+  IN CONST EFI_GUID  *Protocol,
+  IN VOID            *Interface,
+  IN EFI_HANDLE      Handle
+  )
+{
+  //
+  // Disable performance measurement after ExitBootServices because
+  // 1. Performance measurement might impact SMI latency at runtime;
+  // 2. Performance log is copied to non SMRAM at ReadyToBoot so runtime performance
+  //    log is not useful.
+  //
+  mPerformanceMeasurementEnabled = FALSE;
+
+  return EFI_SUCCESS;
+}
 
 /**
   The constructor function initializes the Performance Measurement Enable flag
@@ -40,9 +70,40 @@ SmmPerformanceLibConstructor (
   IN EFI_SYSTEM_TABLE  *SystemTable
   )
 {
+  EFI_STATUS  Status;
+
   mPerformanceMeasurementEnabled =  (BOOLEAN)((PcdGet8 (PcdPerformanceLibraryPropertyMask) & PERFORMANCE_LIBRARY_PROPERTY_MEASUREMENT_ENABLED) != 0);
 
-  return EFI_SUCCESS;
+  Status = gSmst->SmmRegisterProtocolNotify (
+                    &gEdkiiSmmExitBootServicesProtocolGuid,
+                    SmmPerformanceLibExitBootServicesCallback,
+                    &mPerformanceLibExitBootServicesRegistration
+                    );
+  ASSERT_EFI_ERROR (Status);
+
+  return Status;
+}
+
+EFI_STATUS
+EFIAPI
+SmmPerformanceLibDestructor (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_STATUS  Status;
+
+  //
+  // Unregister SmmExitBootServices notification.
+  //
+  Status = gSmst->SmmRegisterProtocolNotify (
+                    &gEdkiiSmmExitBootServicesProtocolGuid,
+                    NULL,
+                    &mPerformanceLibExitBootServicesRegistration
+                    );
+  ASSERT_EFI_ERROR (Status);
+
+  return Status;
 }
 
 /**
diff --git a/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.inf b/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.inf
index d79cd5c8da..002462f5ca 100644
--- a/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.inf
+++ b/MdeModulePkg/Library/SmmPerformanceLib/SmmPerformanceLib.inf
@@ -21,6 +21,7 @@
   LIBRARY_CLASS                  = PerformanceLib|DXE_SMM_DRIVER
 
   CONSTRUCTOR                    = SmmPerformanceLibConstructor
+  DESTRUCTOR                     = SmmPerformanceLibDestructor
 
 #
 # The following information is for reference only and not required by the build tools.
@@ -46,5 +47,8 @@
 [Guids]
   gEdkiiSmmPerformanceMeasurementProtocolGuid          ## SOMETIMES_CONSUMES   ## UNDEFINED # Locate protocol
 
+[Protocols]
+  gEdkiiSmmExitBootServicesProtocolGuid                ## CONSUMES
+
 [Pcd]
   gEfiMdePkgTokenSpaceGuid.PcdPerformanceLibraryPropertyMask    ## CONSUMES
-- 
2.39.1.windows.1