MdeModulePkg DxeCore: Enhance memory profile for memory leak detection

1. Implement include GetRecordingState/SetRecordingState/Record for
memory profile protocol.
2. Consume PcdMemoryProfilePropertyMask to support disable recording
at the start.
3. Consume PcdMemoryProfileDriverPath to control which drivers need
memory profile data.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index e6b9114d2ee5aef564e6a642804d54902b4c69da..743221f6755e343e482218dcf04b3da9523eb29a 100644 (file)
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -2780,11 +2780,13 @@ MemoryProfileInstallProtocol (
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r

-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r


@@ -2795,11 +2797,13 @@ RegisterMemoryProfileImage (
 \r
   @param DriverEntry    Image info.\r
 \r
 \r
   @param DriverEntry    Image info.\r
 \r

-  @retval TRUE          Unregister success.\r
-  @retval FALSE         Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry\r
   );\r
 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry\r
   );\r


@@ -2810,20 +2814,31 @@ UnregisterMemoryProfileImage (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r

+                        EfiMaxMemoryType means the MemoryType is unknown.\r

   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r

+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r

 \r
 \r

-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r

 CoreUpdateProfile (\r
   IN EFI_PHYSICAL_ADDRESS   CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
 CoreUpdateProfile (\r
   IN EFI_PHYSICAL_ADDRESS   CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r

-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r

   );\r
 \r
 /**\r
   );\r
 \r
 /**\r

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf b/MdeModulePkg/Core/Dxe/DxeMain.inf
index e3e4d036dbea8c2cc8a10963010b3ce87c98ee51..450da579e0fc7ec80b2a41a1cc08954b3b1455ef 100644 (file)
--- a/MdeModulePkg/Core/Dxe/DxeMain.inf
+++ b/MdeModulePkg/Core/Dxe/DxeMain.inf
@@ -187,6 +187,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxEfiSystemTablePointerAddress         ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType                 ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask               ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxEfiSystemTablePointerAddress         ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType                 ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask               ## CONSUMES\r

+  gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileDriverPath                 ## CONSUMES\r

   gEfiMdeModulePkgTokenSpaceGuid.PcdPropertiesTableEnable                   ## CONSUMES\r
 \r
 # [Hob]\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdPropertiesTableEnable                   ## CONSUMES\r
 \r
 # [Hob]\r

diff --git a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
index 6626e101592ece266fd1fa613e42a73d93326f16..1ba8488a0e041761c67e3fa5e25cc74684106c01 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
+++ b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
@@ -17,6 +17,9 @@
 \r
 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
 \r
 \r
 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
 \r

+#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \\r
+  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))\r
+\r

 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_CONTEXT        Context;\r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_CONTEXT        Context;\r


@@ -27,12 +30,14 @@ typedef struct {
   UINT32                        Signature;\r
   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;\r
   LIST_ENTRY                    *AllocInfoList;\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;\r
   LIST_ENTRY                    *AllocInfoList;\r

+  CHAR8                         *PdbString;\r

   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_DRIVER_INFO_DATA;\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_DRIVER_INFO_DATA;\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;\r

+  CHAR8                         *ActionString;\r

   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r


@@ -58,7 +63,10 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext
 };\r
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;\r
 \r
 };\r
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;\r
 \r

-BOOLEAN mMemoryProfileRecordingStatus = FALSE;\r
+BOOLEAN mMemoryProfileGettingStatus = FALSE;\r
+BOOLEAN mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+EFI_DEVICE_PATH_PROTOCOL *mMemoryProfileDriverPath;\r
+UINTN                    mMemoryProfileDriverPathSize;\r

 \r
 /**\r
   Get memory profile data.\r
 \r
 /**\r
   Get memory profile data.\r


@@ -69,6 +77,7 @@ BOOLEAN mMemoryProfileRecordingStatus = FALSE;
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r

+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r

   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r
   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r


@@ -90,7 +99,9 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r

-  @return EFI_SUCCESS           Register success.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r

   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
 \r
 **/\r
   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
 \r
 **/\r


@@ -112,7 +123,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r

-  @return EFI_SUCCESS           Unregister success.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r

   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r


@@ -125,10 +138,83 @@ ProfileProtocolUnregisterImage (
   IN UINT64                             ImageSize\r
   );\r
 \r
   IN UINT64                             ImageSize\r
   );\r
 \r

+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  );\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN BOOLEAN                            RecordingState\r
+  );\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolRecord (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  );\r
+\r

 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
   ProfileProtocolGetData,\r
   ProfileProtocolRegisterImage,\r
 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
   ProfileProtocolGetData,\r
   ProfileProtocolRegisterImage,\r

-  ProfileProtocolUnregisterImage\r
+  ProfileProtocolUnregisterImage,\r
+  ProfileProtocolGetRecordingState,\r
+  ProfileProtocolSetRecordingState,\r
+  ProfileProtocolRecord,\r

 };\r
 \r
 /**\r
 };\r
 \r
 /**\r


@@ -307,13 +393,27 @@ BuildDriverInfo (
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   VOID                              *EntryPointInImage;\r
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   VOID                              *EntryPointInImage;\r

+  CHAR8                             *PdbString;\r
+  UINTN                             PdbSize;\r
+  UINTN                             PdbOccupiedSize;\r
+\r
+  PdbSize = 0;\r
+  PdbOccupiedSize = 0;\r
+  PdbString = NULL;\r
+  if (ImageBase != 0) {\r
+    PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);\r
+    if (PdbString != NULL) {\r
+      PdbSize = AsciiStrSize (PdbString);\r
+      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));\r
+    }\r
+  }\r

 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r

-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r

              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r


@@ -325,7 +425,7 @@ BuildDriverInfo (
   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r

-  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+  DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);\r

   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\r
   if (FileName != NULL) {\r
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));\r
   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\r
   if (FileName != NULL) {\r
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));\r


@@ -349,6 +449,14 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage = 0;\r
   DriverInfo->PeakUsage = 0;\r
   DriverInfo->AllocRecordCount = 0;\r
   DriverInfo->CurrentUsage = 0;\r
   DriverInfo->PeakUsage = 0;\r
   DriverInfo->AllocRecordCount = 0;\r

+  if (PdbSize != 0) {\r
+    DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);\r
+    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);\r
+  } else {\r
+    DriverInfo->PdbStringOffset = 0;\r
+    DriverInfoData->PdbString = NULL;\r
+  }\r

 \r
   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
   ContextData->Context.ImageCount ++;\r
 \r
   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
   ContextData->Context.ImageCount ++;\r


@@ -357,6 +465,65 @@ BuildDriverInfo (
   return DriverInfoData;\r
 }\r
 \r
   return DriverInfoData;\r
 }\r
 \r

+/**\r
+  Return if record for this driver is needed..\r
+\r
+  @param DriverFilePath     Driver file path.\r
+\r
+  @retval TRUE              Record for this driver is needed.\r
+  @retval FALSE             Record for this driver is not needed.\r
+\r
+**/\r
+BOOLEAN\r
+NeedRecordThisDriver (\r
+  IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath\r
+  )\r
+{\r
+  EFI_DEVICE_PATH_PROTOCOL          *TmpDevicePath;\r
+  EFI_DEVICE_PATH_PROTOCOL          *DevicePathInstance;\r
+  UINTN                             DevicePathSize;\r
+  UINTN                             FilePathSize;\r
+\r
+  if (!IsDevicePathValid (mMemoryProfileDriverPath, mMemoryProfileDriverPathSize)) {\r
+    //\r
+    // Invalid Device Path means record all.\r
+    //\r
+    return TRUE;\r
+  }\r
+  \r
+  //\r
+  // Record FilePath without END node.\r
+  //\r
+  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);\r
+\r
+  DevicePathInstance = mMemoryProfileDriverPath;\r
+  do {\r
+    //\r
+    // Find END node (it might be END_ENTIRE or END_INSTANCE).\r
+    //\r
+    TmpDevicePath = DevicePathInstance;\r
+    while (!IsDevicePathEndType (TmpDevicePath)) {\r
+      TmpDevicePath = NextDevicePathNode (TmpDevicePath);\r
+    }\r
+\r
+    //\r
+    // Do not compare END node.\r
+    //\r
+    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;\r
+    if ((FilePathSize == DevicePathSize) &&\r
+        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {\r
+      return TRUE;\r
+    }\r
+\r
+    //\r
+    // Get next instance.\r
+    //\r
+    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));\r
+  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);\r
+\r
+  return FALSE;\r
+}\r
+\r

 /**\r
   Register DXE Core to memory profile.\r
 \r
 /**\r
   Register DXE Core to memory profile.\r
 \r


@@ -376,6 +543,8 @@ RegisterDxeCore (
   EFI_PEI_HOB_POINTERS              DxeCoreHob;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   PHYSICAL_ADDRESS                  ImageBase;\r
   EFI_PEI_HOB_POINTERS              DxeCoreHob;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   PHYSICAL_ADDRESS                  ImageBase;\r

+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r

 \r
   ASSERT (ContextData != NULL);\r
 \r
 \r
   ASSERT (ContextData != NULL);\r
 \r


@@ -394,6 +563,14 @@ RegisterDxeCore (
   }\r
   ASSERT (DxeCoreHob.Raw != NULL);\r
 \r
   }\r
   ASSERT (DxeCoreHob.Raw != NULL);\r
 \r

+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DxeCoreHob.MemoryAllocationModule->ModuleName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return FALSE;\r
+  }\r
+\r

   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;\r
   DriverInfoData = BuildDriverInfo (\r
                      ContextData,\r
   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;\r
   DriverInfoData = BuildDriverInfo (\r
                      ContextData,\r


@@ -433,7 +610,14 @@ MemoryProfileInit (
     return;\r
   }\r
 \r
     return;\r
   }\r
 \r

-  mMemoryProfileRecordingStatus = TRUE;\r
+  mMemoryProfileGettingStatus = FALSE;\r
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  } else {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;\r
+  }\r
+  mMemoryProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);\r
+  mMemoryProfileDriverPath = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));\r

   mMemoryProfileContextPtr = &mMemoryProfileContext;\r
 \r
   RegisterDxeCore (HobStart, &mMemoryProfileContext);\r
   mMemoryProfileContextPtr = &mMemoryProfileContext;\r
 \r
   RegisterDxeCore (HobStart, &mMemoryProfileContext);\r


@@ -504,11 +688,13 @@ GetFileNameFromFilePath (
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r

-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r


@@ -518,12 +704,16 @@ RegisterMemoryProfileImage (
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r


@@ -536,10 +726,10 @@ RegisterMemoryProfileImage (
                      FileType\r
                      );\r
   if (DriverInfoData == NULL) {\r
                      FileType\r
                      );\r
   if (DriverInfoData == NULL) {\r

-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r

   }\r
 \r
   }\r
 \r

-  return TRUE;\r
+  return EFI_SUCCESS;\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r


@@ -586,45 +776,9 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   return NULL;\r
 }\r
 \r
   return NULL;\r
 }\r
 \r

-/**\r
-  Search dummy image from memory profile.\r
-\r
-  @param ContextData    Memory profile context.\r
-\r
-  @return Pointer to memory profile driver info.\r
-\r
-**/\r
-MEMORY_PROFILE_DRIVER_INFO_DATA *\r
-FindDummyImage (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData\r
-  )\r
-{\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
-\r
-  DriverInfoList = ContextData->DriverInfoList;\r
-\r
-  for (DriverLink = DriverInfoList->ForwardLink;\r
-       DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
-    DriverInfoData = CR (\r
-                   DriverLink,\r
-                   MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                   Link,\r
-                   MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                   );\r
-    if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {\r
-      return DriverInfoData;\r
-    }\r
-  }\r
-\r
-  return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);\r
-}\r
-\r

 /**\r
   Search image from memory profile.\r
 /**\r
   Search image from memory profile.\r

-  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)\r
+  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize).\r

 \r
   @param ContextData    Memory profile context.\r
   @param Address        Image or Function address.\r
 \r
   @param ContextData    Memory profile context.\r
   @param Address        Image or Function address.\r


@@ -661,10 +815,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
     }\r
   }\r
 \r

-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r


@@ -672,11 +823,13 @@ GetMemoryProfileDriverInfoFromAddress (
 \r
   @param DriverEntry    Image info.\r
 \r
 \r
   @param DriverEntry    Image info.\r
 \r

-  @retval TRUE          Unregister success.\r
-  @retval FALSE         Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry\r
   )\r
 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry\r
   )\r


@@ -689,12 +842,16 @@ UnregisterMemoryProfileImage (
   VOID                              *EntryPointInImage;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
   VOID                              *EntryPointInImage;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   DriverInfoData = NULL;\r
   }\r
 \r
   DriverInfoData = NULL;\r


@@ -716,12 +873,13 @@ UnregisterMemoryProfileImage (
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\r
   if (DriverInfoData == NULL) {\r
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\r
   if (DriverInfoData == NULL) {\r

-    return FALSE;\r
+    return EFI_NOT_FOUND;\r

   }\r
 \r
   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
 \r
   }\r
 \r
   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
 \r

-  DriverInfoData->DriverInfo.ImageBase = 0;\r
+  // Keep the ImageBase for RVA calculation in Application.\r
+  //DriverInfoData->DriverInfo.ImageBase = 0;\r

   DriverInfoData->DriverInfo.ImageSize = 0;\r
 \r
   if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r
   DriverInfoData->DriverInfo.ImageSize = 0;\r
 \r
   if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r


@@ -733,7 +891,7 @@ UnregisterMemoryProfileImage (
     CoreInternalFreePool (DriverInfoData, NULL);\r
   }\r
 \r
     CoreInternalFreePool (DriverInfoData, NULL);\r
   }\r
 \r

-  return TRUE;\r
+  return EFI_SUCCESS;\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r


@@ -803,55 +961,80 @@ GetProfileMemoryIndex (
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r

+  @param ActionString   String for memory profile action.\r

 \r
 \r

-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 CoreUpdateProfileAllocate (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,\r
 CoreUpdateProfileAllocate (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,\r

-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r

   )\r
 {\r
   EFI_STATUS                        Status;\r
   )\r
 {\r
   EFI_STATUS                        Status;\r

-  MEMORY_PROFILE_CONTEXT           *Context;\r
-  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\r
-  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;\r
+  MEMORY_PROFILE_CONTEXT            *Context;\r
+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
+  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;\r

   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   UINTN                             ProfileMemoryIndex;\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   UINTN                             ProfileMemoryIndex;\r

+  MEMORY_PROFILE_ACTION             BasicAction;\r
+  UINTN                             ActionStringSize;\r
+  UINTN                             ActionStringOccupiedSize;\r

 \r
 \r

-  AllocInfoData = NULL;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r

 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r

-  ASSERT (DriverInfoData != NULL);\r
+  if (DriverInfoData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  ActionStringSize = 0;\r
+  ActionStringOccupiedSize = 0;\r
+  if (ActionString != NULL) {\r
+    ActionStringSize = AsciiStrSize (ActionString);\r
+    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));\r
+  }\r

 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r

+  AllocInfoData = NULL;\r

   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r

-             sizeof (*AllocInfoData),\r
+             sizeof (*AllocInfoData) + ActionStringSize,\r

              (VOID **) &AllocInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
              (VOID **) &AllocInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r

-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r

   }\r
   ASSERT (AllocInfoData != NULL);\r
   }\r
   ASSERT (AllocInfoData != NULL);\r

+\r
+  //\r
+  // Only update SequenceCount if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ContextData->Context.SequenceCount ++;\r
+  }\r
+\r

   AllocInfo = &AllocInfoData->AllocInfo;\r
   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo = &AllocInfoData->AllocInfo;\r
   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r

-  AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+  AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);\r

   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
   AllocInfo->CallerAddress      = CallerAddress;\r
   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r
   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
   AllocInfo->CallerAddress      = CallerAddress;\r
   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r


@@ -859,50 +1042,64 @@ CoreUpdateProfileAllocate (
   AllocInfo->MemoryType         = MemoryType;\r
   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
   AllocInfo->Size               = Size;\r
   AllocInfo->MemoryType         = MemoryType;\r
   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
   AllocInfo->Size               = Size;\r

+  if (ActionString != NULL) {\r
+    AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+    AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);\r
+    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);\r
+  } else {\r
+    AllocInfo->ActionStringOffset = 0;\r
+    AllocInfoData->ActionString = NULL;\r
+  }\r

 \r
   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\r
 \r
 \r
   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\r
 \r

-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
-\r
+  Context = &ContextData->Context;\r

   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfo = &DriverInfoData->DriverInfo;\r

-  DriverInfo->CurrentUsage += Size;\r
-  if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
-    DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
-  }\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
-  if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
-    DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
-  }\r

   DriverInfo->AllocRecordCount ++;\r
 \r
   DriverInfo->AllocRecordCount ++;\r
 \r

-  Context = &ContextData->Context;\r
-  Context->CurrentTotalUsage += Size;\r
-  if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
-    Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
-  }\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
-  if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
-    Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+  //\r
+  // Update summary if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+\r
+    DriverInfo->CurrentUsage += Size;\r
+    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
+      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
+    }\r
+    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
+    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
+      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    Context->CurrentTotalUsage += Size;\r
+    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
+      Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
+    }\r
+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+    }\r

   }\r
   }\r

-  Context->SequenceCount ++;\r

 \r
 \r

-  return TRUE;\r
+  return EFI_SUCCESS;\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r

-  Get memory profile alloc info from memory profile\r
+  Get memory profile alloc info from memory profile.\r

 \r
 \r

-  @param DriverInfoData     Driver info\r
-  @param Action             This Free action\r
-  @param Size               Buffer size\r
-  @param Buffer             Buffer address\r
+  @param DriverInfoData     Driver info.\r
+  @param BasicAction        This Free basic action.\r
+  @param Size               Buffer size.\r
+  @param Buffer             Buffer address.\r

 \r
   @return Pointer to memory profile alloc info.\r
 \r
   @return Pointer to memory profile alloc info.\r

+\r

 **/\r
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r
 **/\r
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r

-  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN MEMORY_PROFILE_ACTION              BasicAction,\r

   IN UINTN                              Size,\r
   IN VOID                               *Buffer\r
   )\r
   IN UINTN                              Size,\r
   IN VOID                               *Buffer\r
   )\r


@@ -924,10 +1121,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                       );\r
     AllocInfo = &AllocInfoData->AllocInfo;\r
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                       );\r
     AllocInfo = &AllocInfoData->AllocInfo;\r

-    if (AllocInfo->Action != Action) {\r
+    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {\r

       continue;\r
     }\r
       continue;\r
     }\r

-    switch (Action) {\r
+    switch (BasicAction) {\r

       case MemoryProfileActionAllocatePages:\r
         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r
       case MemoryProfileActionAllocatePages:\r
         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r


@@ -956,11 +1153,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
 \r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
 \r

-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r

 CoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
 CoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r


@@ -978,114 +1176,140 @@ CoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   UINTN                            ProfileMemoryIndex;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   UINTN                            ProfileMemoryIndex;\r

+  MEMORY_PROFILE_ACTION            BasicAction;\r
+  BOOLEAN                          Found;\r
+\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r

 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r

-  ASSERT (DriverInfoData != NULL);\r

 \r
 \r

-  switch (Action) {\r
-    case MemoryProfileActionFreePages:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
-      break;\r
-    case MemoryProfileActionFreePool:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
-      break;\r
-    default:\r
-      ASSERT (FALSE);\r
-      AllocInfoData = NULL;\r
-      break;\r
-  }\r
-  if (AllocInfoData == NULL) {\r
-    //\r
-    // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
-    //\r
-    DriverInfoList = ContextData->DriverInfoList;\r
-\r
-    for (DriverLink = DriverInfoList->ForwardLink;\r
-         DriverLink != DriverInfoList;\r
-         DriverLink = DriverLink->ForwardLink) {\r
-      ThisDriverInfoData = CR (\r
-                             DriverLink,\r
-                             MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                             Link,\r
-                             MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                             );\r
-      switch (Action) {\r
+  //\r
+  // Do not return if DriverInfoData == NULL here,\r
+  // because driver A might free memory allocated by driver B.\r
+  //\r
+\r
+  //\r
+  // Need use do-while loop to find all possible records,\r
+  // because one address might be recorded multiple times.\r
+  //\r
+  Found = FALSE;\r
+  AllocInfoData = NULL;\r
+  do {\r
+    if (DriverInfoData != NULL) {\r
+      switch (BasicAction) {\r

         case MemoryProfileActionFreePages:\r
         case MemoryProfileActionFreePages:\r

-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r

           break;\r
         case MemoryProfileActionFreePool:\r
           break;\r
         case MemoryProfileActionFreePool:\r

-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r

           break;\r
         default:\r
           ASSERT (FALSE);\r
           AllocInfoData = NULL;\r
           break;\r
       }\r
           break;\r
         default:\r
           ASSERT (FALSE);\r
           AllocInfoData = NULL;\r
           break;\r
       }\r

-      if (AllocInfoData != NULL) {\r
-        DriverInfoData = ThisDriverInfoData;\r
-        break;\r
-      }\r

     }\r
     }\r

-\r

     if (AllocInfoData == NULL) {\r
       //\r
     if (AllocInfoData == NULL) {\r
       //\r

-      // No matched allocate operation is found for this free operation.\r
-      // It is because the specified memory type allocate operation has been\r
-      // filtered by CoreNeedRecordProfile(), but free operations have no\r
-      // memory type information, they can not be filtered by CoreNeedRecordProfile().\r
-      // Then, they will be filtered here.\r
+      // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r

       //\r
       //\r

-      return FALSE;\r
-    }\r
-  }\r
+      DriverInfoList = ContextData->DriverInfoList;\r
+\r
+      for (DriverLink = DriverInfoList->ForwardLink;\r
+           DriverLink != DriverInfoList;\r
+           DriverLink = DriverLink->ForwardLink) {\r
+        ThisDriverInfoData = CR (\r
+                               DriverLink,\r
+                               MEMORY_PROFILE_DRIVER_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                               );\r
+        switch (BasicAction) {\r
+          case MemoryProfileActionFreePages:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+            break;\r
+          case MemoryProfileActionFreePool:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+            break;\r
+          default:\r
+            ASSERT (FALSE);\r
+            AllocInfoData = NULL;\r
+            break;\r
+        }\r
+        if (AllocInfoData != NULL) {\r
+          DriverInfoData = ThisDriverInfoData;\r
+          break;\r
+        }\r
+      }\r

 \r
 \r

-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\r
+      if (AllocInfoData == NULL) {\r
+        //\r
+        // If (!Found), no matched allocate info is found for this free action.\r
+        // It is because the specified memory type allocate actions have been filtered by\r
+        // CoreNeedRecordProfile(), but free actions may have no memory type information,\r
+        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be\r
+        // filtered here.\r
+        //\r
+        // If (Found), it is normal exit path.\r
+        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);\r
+      }\r
+    }\r

 \r
 \r

-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+    Found = TRUE;\r

 \r
 \r

-  Context->CurrentTotalUsage -= AllocInfo->Size;\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+    Context = &ContextData->Context;\r
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    AllocInfo = &AllocInfoData->AllocInfo;\r

 \r
 \r

-  DriverInfo->CurrentUsage -= AllocInfo->Size;\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
-  DriverInfo->AllocRecordCount --;\r
+    DriverInfo->AllocRecordCount --;\r
+    //\r
+    // Update summary if and only if it is basic action.\r
+    //\r
+    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {\r
+      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r

 \r
 \r

-  RemoveEntryList (&AllocInfoData->Link);\r
+      Context->CurrentTotalUsage -= AllocInfo->Size;\r
+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r

 \r
 \r

-  if (Action == MemoryProfileActionFreePages) {\r
-    if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
-      CoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
-        (VOID *) (UINTN) AllocInfo->Buffer\r
-        );\r
+      DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r

     }\r
     }\r

-    if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
-      CoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
-        (VOID *) ((UINTN) Buffer + Size)\r
-        );\r
-    }\r
-  }\r

 \r
 \r

-  //\r
-  // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  CoreInternalFreePool (AllocInfoData, NULL);\r
+    RemoveEntryList (&AllocInfoData->Link);\r
+\r
+    if (BasicAction == MemoryProfileActionFreePages) {\r
+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+        CoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
+          (VOID *) (UINTN) AllocInfo->Buffer,\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
+        CoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
+          (VOID *) ((UINTN) Buffer + Size),\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+    }\r

 \r
 \r

-  return TRUE;\r
+    //\r
+    // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    CoreInternalFreePool (AllocInfoData, NULL);\r
+  } while (TRUE);\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r


@@ -1094,62 +1318,90 @@ CoreUpdateProfileFree (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r

+                        EfiMaxMemoryType means the MemoryType is unknown.\r

   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r

-\r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r

 \r
 **/\r
 \r
 **/\r

-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r

 CoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
 CoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r

-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r

   )\r
 {\r
   )\r
 {\r

+  EFI_STATUS                    Status;\r

   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r

+  MEMORY_PROFILE_ACTION         BasicAction;\r

 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   }\r
 \r

-  if (!mMemoryProfileRecordingStatus) {\r
-    return FALSE;\r
+  if (mMemoryProfileGettingStatus) {\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
+  if (!mMemoryProfileRecordingEnable) {\r
+    return EFI_ABORTED;\r

   }\r
 \r
   //\r
   }\r
 \r
   //\r

-  // Only record limited MemoryType.\r
+  // Get the basic action to know how to process the record\r

   //\r
   //\r

-  if (!CoreNeedRecordProfile (MemoryType)) {\r
-    return FALSE;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
+  //\r
+  // EfiMaxMemoryType means the MemoryType is unknown.\r
+  //\r
+  if (MemoryType != EfiMaxMemoryType) {\r
+    //\r
+    // Only record limited MemoryType.\r
+    //\r
+    if (!CoreNeedRecordProfile (MemoryType)) {\r
+      return EFI_UNSUPPORTED;\r
+    }\r

   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r

-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r

   }\r
 \r
   }\r
 \r

-  switch (Action) {\r
+  switch (BasicAction) {\r

     case MemoryProfileActionAllocatePages:\r
     case MemoryProfileActionAllocatePages:\r

-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r

       break;\r
     case MemoryProfileActionFreePages:\r
       break;\r
     case MemoryProfileActionFreePages:\r

-      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r

       break;\r
     case MemoryProfileActionAllocatePool:\r
       break;\r
     case MemoryProfileActionAllocatePool:\r

-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r

       break;\r
     case MemoryProfileActionFreePool:\r
       break;\r
     case MemoryProfileActionFreePool:\r

-      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r

       break;\r
     default:\r
       ASSERT (FALSE);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r

+      Status = EFI_UNSUPPORTED;\r

       break;\r
   }\r
       break;\r
   }\r

-  return TRUE;\r
+  return Status;\r

 }\r
 \r
 ////////////////////\r
 }\r
 \r
 ////////////////////\r


@@ -1167,8 +1419,11 @@ MemoryProfileGetDataSize (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
 {\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r

+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r

   LIST_ENTRY                        *DriverInfoList;\r
   LIST_ENTRY                        *DriverLink;\r
   LIST_ENTRY                        *DriverInfoList;\r
   LIST_ENTRY                        *DriverLink;\r

+  LIST_ENTRY                        *AllocInfoList;\r
+  LIST_ENTRY                        *AllocLink;\r

   UINTN                             TotalSize;\r
 \r
 \r
   UINTN                             TotalSize;\r
 \r
 \r


@@ -1178,7 +1433,6 @@ MemoryProfileGetDataSize (
   }\r
 \r
   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);\r
   }\r
 \r
   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);\r

-  TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;\r

 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r


@@ -1190,7 +1444,20 @@ MemoryProfileGetDataSize (
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r

-    TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;\r
+    TotalSize += DriverInfoData->DriverInfo.Header.Length;\r
+\r
+    AllocInfoList = DriverInfoData->AllocInfoList;\r
+    for (AllocLink = AllocInfoList->ForwardLink;\r
+         AllocLink != AllocInfoList;\r
+         AllocLink = AllocLink->ForwardLink) {\r
+      AllocInfoData = CR (\r
+                        AllocLink,\r
+                        MEMORY_PROFILE_ALLOC_INFO_DATA,\r
+                        Link,\r
+                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
+                        );\r
+      TotalSize += AllocInfoData->AllocInfo.Header.Length;\r
+    }\r

   }\r
 \r
   return TotalSize;\r
   }\r
 \r
   return TotalSize;\r


@@ -1217,6 +1484,8 @@ MemoryProfileCopyData (
   LIST_ENTRY                        *DriverLink;\r
   LIST_ENTRY                        *AllocInfoList;\r
   LIST_ENTRY                        *AllocLink;\r
   LIST_ENTRY                        *DriverLink;\r
   LIST_ENTRY                        *AllocInfoList;\r
   LIST_ENTRY                        *AllocLink;\r

+  UINTN                             PdbSize;\r
+  UINTN                             ActionStringSize;\r

 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r


@@ -1238,7 +1507,11 @@ MemoryProfileCopyData (
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r

-    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);\r
+    if (DriverInfo->PdbStringOffset != 0) {\r
+      PdbSize = AsciiStrSize (DriverInfoData->PdbString);\r
+      CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);\r
+    }\r
+    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);\r

 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r


@@ -1251,10 +1524,14 @@ MemoryProfileCopyData (
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r

-      AllocInfo += 1;\r
+      if (AllocInfo->ActionStringOffset != 0) {\r
+        ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);\r
+        CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);\r
+      }\r
+      AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) AllocInfo + AllocInfo->Header.Length);\r

     }\r
 \r
     }\r
 \r

-    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);\r
+    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)  AllocInfo;\r

   }\r
 }\r
 \r
   }\r
 }\r
 \r


@@ -1267,6 +1544,7 @@ MemoryProfileCopyData (
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r

+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r

   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r
   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r


@@ -1281,28 +1559,28 @@ ProfileProtocolGetData (
 {\r
   UINTN                                 Size;\r
   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
 {\r
   UINTN                                 Size;\r
   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r

-  BOOLEAN                               MemoryProfileRecordingStatus;\r
+  BOOLEAN                               MemoryProfileGettingStatus;\r

 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r

-  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;\r
-  mMemoryProfileRecordingStatus = FALSE;\r
+  MemoryProfileGettingStatus = mMemoryProfileGettingStatus;\r
+  mMemoryProfileGettingStatus = TRUE;\r

 \r
   Size = MemoryProfileGetDataSize ();\r
 \r
   if (*ProfileSize < Size) {\r
     *ProfileSize = Size;\r
 \r
   Size = MemoryProfileGetDataSize ();\r
 \r
   if (*ProfileSize < Size) {\r
     *ProfileSize = Size;\r

-    mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+    mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r

     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
   *ProfileSize = Size;\r
   MemoryProfileCopyData (ProfileBuffer);\r
 \r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
   *ProfileSize = Size;\r
   MemoryProfileCopyData (ProfileBuffer);\r
 \r

-  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+  mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r

   return EFI_SUCCESS;\r
 }\r
 \r
   return EFI_SUCCESS;\r
 }\r
 \r


@@ -1315,8 +1593,10 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r

-  @return EFI_SUCCESS           Register success.\r
-  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r

 \r
 **/\r
 EFI_STATUS\r
 \r
 **/\r
 EFI_STATUS\r


@@ -1342,7 +1622,7 @@ ProfileProtocolRegisterImage (
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);\r
 \r
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);\r
 \r

-  return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;\r
+  return RegisterMemoryProfileImage (&DriverEntry, FileType);\r

 }\r
 \r
 /**\r
 }\r
 \r
 /**\r


@@ -1353,7 +1633,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r

-  @return EFI_SUCCESS           Unregister success.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r

   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r


@@ -1378,7 +1660,105 @@ ProfileProtocolUnregisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r

-  return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;\r
+  return UnregisterMemoryProfileImage (&DriverEntry);\r
+}\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (RecordingState == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  *RecordingState = mMemoryProfileRecordingEnable;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN BOOLEAN                            RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  mMemoryProfileRecordingEnable = RecordingState;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolRecord (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  )\r
+{\r
+  return CoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r

 }\r
 \r
 ////////////////////\r
 }\r
 \r
 ////////////////////\r

diff --git a/MdeModulePkg/Core/Dxe/Mem/Page.c b/MdeModulePkg/Core/Dxe/Mem/Page.c
index 898b722a431989d94b6816a23387d5204bb8aa3d..b02bafb4bed9de2e66573f91f2290e7cc796caa0 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/Page.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Page.c
@@ -1335,7 +1335,14 @@ CoreAllocatePages (
 \r
   Status = CoreInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);\r
   if (!EFI_ERROR (Status)) {\r
 \r
   Status = CoreInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);\r
   if (!EFI_ERROR (Status)) {\r

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) *Memory,\r
+      NULL\r
+      );\r

     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r
     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r


@@ -1444,7 +1451,14 @@ CoreFreePages (
 \r
   Status = CoreInternalFreePages (Memory, NumberOfPages, &MemoryType);\r
   if (!EFI_ERROR (Status)) {\r
 \r
   Status = CoreInternalFreePages (Memory, NumberOfPages, &MemoryType);\r
   if (!EFI_ERROR (Status)) {\r

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) Memory,\r
+      NULL\r
+      );\r

     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r
     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r

diff --git a/MdeModulePkg/Core/Dxe/Mem/Pool.c b/MdeModulePkg/Core/Dxe/Mem/Pool.c
index 2980e2293e3d1dadd59e4130cb6c7387aef8beda..6ef5fba395f5805794d3972746f61fb365e66c45 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/Pool.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Pool.c
@@ -276,7 +276,14 @@ CoreAllocatePool (
 \r
   Status = CoreInternalAllocatePool (PoolType, Size, Buffer);\r
   if (!EFI_ERROR (Status)) {\r
 \r
   Status = CoreInternalAllocatePool (PoolType, Size, Buffer);\r
   if (!EFI_ERROR (Status)) {\r

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePool,\r
+      PoolType,\r
+      Size,\r
+      *Buffer,\r
+      NULL\r
+      );\r

     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r
     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r


@@ -505,7 +512,14 @@ CoreFreePool (
 \r
   Status = CoreInternalFreePool (Buffer, &PoolType);\r
   if (!EFI_ERROR (Status)) {\r
 \r
   Status = CoreInternalFreePool (Buffer, &PoolType);\r
   if (!EFI_ERROR (Status)) {\r

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, PoolType, 0, Buffer);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePool,\r
+      PoolType,\r
+      0,\r
+      Buffer,\r
+      NULL\r
+      );\r

     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r
     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r