]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
MdeModulePkg DxeCore: Remove unreferenced symbol for memory profile
[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / MemoryProfileRecord.c
index 1ae991b7024e62e7906fb48738163a470bc0c9b5..a91a719b4de5593b1d21f0aaf0951ace256e5f5e 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   Support routines for UEFI memory profile.\r
 \r
-  Copyright (c) 2014 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2014 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
 **/\r
 \r
 #include "DxeMain.h"\r
+#include "Imem.h"\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
@@ -26,12 +30,14 @@ typedef struct {
   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
+  CHAR8                         *ActionString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
@@ -57,7 +63,11 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext
 };\r
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;\r
 \r
-BOOLEAN mMemoryProfileRecordingStatus = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_LOCK mMemoryProfileLock = EFI_INITIALIZE_LOCK_VARIABLE (TPL_NOTIFY);\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN mMemoryProfileGettingStatus = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_DEVICE_PATH_PROTOCOL *mMemoryProfileDriverPath;\r
+GLOBAL_REMOVE_IF_UNREFERENCED UINTN                    mMemoryProfileDriverPathSize;\r
 \r
 /**\r
   Get memory profile data.\r
@@ -68,6 +78,7 @@ BOOLEAN mMemoryProfileRecordingStatus = FALSE;
   @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
@@ -89,7 +100,9 @@ ProfileProtocolGetData (
   @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
@@ -111,7 +124,9 @@ ProfileProtocolRegisterImage (
   @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
@@ -124,12 +139,107 @@ ProfileProtocolUnregisterImage (
   IN UINT64                             ImageSize\r
   );\r
 \r
-EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\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
+GLOBAL_REMOVE_IF_UNREFERENCED EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
   ProfileProtocolGetData,\r
   ProfileProtocolRegisterImage,\r
-  ProfileProtocolUnregisterImage\r
+  ProfileProtocolUnregisterImage,\r
+  ProfileProtocolGetRecordingState,\r
+  ProfileProtocolSetRecordingState,\r
+  ProfileProtocolRecord,\r
 };\r
 \r
+/**\r
+  Acquire lock on mMemoryProfileLock.\r
+**/\r
+VOID\r
+CoreAcquireMemoryProfileLock (\r
+  VOID\r
+  )\r
+{\r
+  CoreAcquireLock (&mMemoryProfileLock);\r
+}\r
+\r
+/**\r
+  Release lock on mMemoryProfileLock.\r
+**/\r
+VOID\r
+CoreReleaseMemoryProfileLock (\r
+  VOID\r
+  )\r
+{\r
+  CoreReleaseLock (&mMemoryProfileLock);\r
+}\r
+\r
 /**\r
   Return memory profile context.\r
 \r
@@ -306,25 +416,40 @@ BuildDriverInfo (
   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
-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r
              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
     return NULL;\r
   }\r
+  ASSERT (DriverInfoData != NULL);\r
 \r
   ZeroMem (DriverInfoData, sizeof (*DriverInfoData));\r
 \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
@@ -348,6 +473,14 @@ BuildDriverInfo (
   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
@@ -356,6 +489,65 @@ BuildDriverInfo (
   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
@@ -375,6 +567,8 @@ RegisterDxeCore (
   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
@@ -393,6 +587,14 @@ RegisterDxeCore (
   }\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
@@ -432,7 +634,14 @@ MemoryProfileInit (
     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
@@ -503,11 +712,13 @@ GetFileNameFromFilePath (
   @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
-BOOLEAN\r
+EFI_STATUS\r
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
@@ -517,12 +728,16 @@ RegisterMemoryProfileImage (
   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
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
@@ -535,10 +750,10 @@ RegisterMemoryProfileImage (
                      FileType\r
                      );\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -585,45 +800,9 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   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
-  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
@@ -660,10 +839,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -671,11 +847,13 @@ GetMemoryProfileDriverInfoFromAddress (
 \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
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry\r
   )\r
@@ -688,12 +866,16 @@ UnregisterMemoryProfileImage (
   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
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = NULL;\r
@@ -715,12 +897,13 @@ UnregisterMemoryProfileImage (
     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
-  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
@@ -729,16 +912,17 @@ UnregisterMemoryProfileImage (
     //\r
     // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
     //\r
-    CoreInternalFreePool (DriverInfoData);\r
+    CoreInternalFreePool (DriverInfoData, NULL);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
   Return if this memory type needs to be recorded into memory profile.\r
-  If BIOS memory type (0 ~ EfiMaxMemoryType), it checks bit (1 << MemoryType).\r
+  If BIOS memory type (0 ~ EfiMaxMemoryType - 1), it checks bit (1 << MemoryType).\r
   If OS memory type (0x80000000 ~ 0xFFFFFFFF), it checks bit63 - 0x8000000000000000.\r
+  If OEM memory type (0x70000000 ~ 0x7FFFFFFF), it checks bit62 - 0x4000000000000000.\r
 \r
   @param MemoryType     Memory type.\r
 \r
@@ -753,8 +937,10 @@ CoreNeedRecordProfile (
 {\r
   UINT64 TestBit;\r
 \r
-  if ((UINT32) MemoryType >= 0x80000000) {\r
+  if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
     TestBit = BIT63;\r
+  } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
+    TestBit = BIT62;\r
   } else {\r
     TestBit = LShiftU64 (1, MemoryType);\r
   }\r
@@ -768,21 +954,24 @@ CoreNeedRecordProfile (
 \r
 /**\r
   Convert EFI memory type to profile memory index. The rule is:\r
-  If BIOS memory type (0 ~ EfiMaxMemoryType), ProfileMemoryIndex = MemoryType.\r
+  If BIOS memory type (0 ~ EfiMaxMemoryType - 1), ProfileMemoryIndex = MemoryType.\r
   If OS memory type (0x80000000 ~ 0xFFFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType.\r
+  If OEM memory type (0x70000000 ~ 0x7FFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType + 1.\r
 \r
   @param MemoryType     Memory type.\r
 \r
-  @return EFI memory type as profile memory index.\r
+  @return Profile memory index.\r
 \r
 **/\r
-EFI_MEMORY_TYPE\r
+UINTN\r
 GetProfileMemoryIndex (\r
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
 {\r
-  if ((UINT32) MemoryType >= 0x80000000) {\r
+  if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
     return EfiMaxMemoryType;\r
+  } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
+    return EfiMaxMemoryType + 1;\r
   } else {\r
     return MemoryType;\r
   }\r
@@ -796,55 +985,80 @@ GetProfileMemoryIndex (
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
+  @param ActionString   String for memory profile action.\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
-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
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\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
-  EFI_MEMORY_TYPE                   ProfileMemoryIndex;\r
+  UINTN                             ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION             BasicAction;\r
+  UINTN                             ActionStringSize;\r
+  UINTN                             ActionStringOccupiedSize;\r
 \r
-  AllocInfoData = NULL;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\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
+  AllocInfoData = NULL;\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
-             sizeof (*AllocInfoData),\r
+             sizeof (*AllocInfoData) + ActionStringSize,\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
+  //\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->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
@@ -852,50 +1066,64 @@ CoreUpdateProfileAllocate (
   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
-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
-\r
+  Context = &ContextData->Context;\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
-  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
-  Context->SequenceCount ++;\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
-  Get memory profile alloc info from memory profile\r
+  Get memory profile alloc info from memory profile.\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
 **/\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
@@ -917,10 +1145,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       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
-    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
@@ -949,11 +1177,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @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
-BOOLEAN\r
+EFI_STATUS\r
 CoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
@@ -970,115 +1199,144 @@ CoreUpdateProfileFree (
   LIST_ENTRY                       *DriverInfoList;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
-  EFI_MEMORY_TYPE                  ProfileMemoryIndex;\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
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\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
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\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
-      if (AllocInfoData != NULL) {\r
-        DriverInfoData = ThisDriverInfoData;\r
-        break;\r
-      }\r
     }\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
-      return FALSE;\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
+      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
-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\r
+    ASSERT (DriverInfoData != NULL);\r
+    ASSERT (AllocInfoData != NULL);\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+    Found = TRUE;\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
-  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
-  RemoveEntryList (&AllocInfoData->Link);\r
+      Context->CurrentTotalUsage -= AllocInfo->Size;\r
+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\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
-    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
-  // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  CoreInternalFreePool (AllocInfoData);\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
-  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
@@ -1087,67 +1345,93 @@ 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
+                        EfiMaxMemoryType means the MemoryType is unknown.\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
-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, // Valid for AllocatePages/AllocatePool\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
+  EFI_STATUS                    Status;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  MEMORY_PROFILE_ACTION         BasicAction;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\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
+  // Get the basic action to know how to process the record\r
   //\r
-  // Free operations have no memory type information, so skip the check.\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
+  //\r
+  // EfiMaxMemoryType means the MemoryType is unknown.\r
   //\r
-  if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {\r
+  if (MemoryType != EfiMaxMemoryType) {\r
     //\r
     // Only record limited MemoryType.\r
     //\r
     if (!CoreNeedRecordProfile (MemoryType)) {\r
-      return FALSE;\r
+      return EFI_UNSUPPORTED;\r
     }\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  switch (Action) {\r
+  CoreAcquireMemoryProfileLock ();\r
+  switch (BasicAction) {\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
-      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\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
-      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r
+      Status = EFI_UNSUPPORTED;\r
       break;\r
   }\r
-  return TRUE;\r
+  CoreReleaseMemoryProfileLock ();\r
+\r
+  return Status;\r
 }\r
 \r
 ////////////////////\r
@@ -1165,8 +1449,11 @@ MemoryProfileGetDataSize (
 {\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                        *AllocInfoList;\r
+  LIST_ENTRY                        *AllocLink;\r
   UINTN                             TotalSize;\r
 \r
 \r
@@ -1176,7 +1463,6 @@ MemoryProfileGetDataSize (
   }\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
@@ -1188,7 +1474,20 @@ MemoryProfileGetDataSize (
                        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
@@ -1215,6 +1514,8 @@ MemoryProfileCopyData (
   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
@@ -1236,7 +1537,11 @@ MemoryProfileCopyData (
                        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
@@ -1249,10 +1554,14 @@ MemoryProfileCopyData (
                         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
-    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
@@ -1265,6 +1574,7 @@ MemoryProfileCopyData (
   @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
@@ -1279,28 +1589,28 @@ ProfileProtocolGetData (
 {\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
-  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;\r
-  mMemoryProfileRecordingStatus = FALSE;\r
+  MemoryProfileGettingStatus = mMemoryProfileGettingStatus;\r
+  mMemoryProfileGettingStatus = TRUE;\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
-  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+  mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1313,8 +1623,10 @@ ProfileProtocolGetData (
   @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
@@ -1340,7 +1652,7 @@ ProfileProtocolRegisterImage (
   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
@@ -1351,7 +1663,9 @@ ProfileProtocolRegisterImage (
   @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
@@ -1376,7 +1690,105 @@ ProfileProtocolUnregisterImage (
   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