UefiCpuPkg: Enhance logic in InitializeMpExceptionStackSwitchHandlers

[mirror_edk2.git] / UefiCpuPkg / CpuDxe / CpuMp.c

diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/UefiCpuPkg/CpuDxe/CpuMp.c
index f3ca813d2a1ce19c875dffcc8aa82c75b33df6d2..e7575d9b8062e77b4a1957338d32d9679d258cca 100644 (file)
--- a/UefiCpuPkg/CpuDxe/CpuMp.c
+++ b/UefiCpuPkg/CpuDxe/CpuMp.c
@@ -600,8 +600,9 @@ CollectBistDataFromHob (
 // Structure for InitializeSeparateExceptionStacks\r
 //\r
 typedef struct {\r
-  VOID     *Buffer;\r
-  UINTN    *BufferSize;\r
+  VOID          *Buffer;\r
+  UINTN         BufferSize;\r
+  EFI_STATUS    Status;\r
 } EXCEPTION_STACK_SWITCH_CONTEXT;\r
 \r
 /**\r
@@ -620,9 +621,18 @@ InitializeExceptionStackSwitchHandlers (
   )\r
 {\r
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;\r
+  UINTN                           Index;\r
 \r
+  MpInitLibWhoAmI (&Index);\r
   SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;\r
-  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->BufferSize);\r
+\r
+  //\r
+  // This may be called twice for each Cpu. Only run InitializeSeparateExceptionStacks\r
+  // if this is the first call or the first call failed because of size too small.\r
+  //\r
+  if ((SwitchStackData[Index].Status == EFI_NOT_STARTED) || (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL)) {\r
+    SwitchStackData[Index].Status = InitializeSeparateExceptionStacks (SwitchStackData[Index].Buffer, &SwitchStackData[Index].BufferSize);\r
+  }\r
 }\r
 \r
 /**\r
@@ -638,53 +648,69 @@ InitializeMpExceptionStackSwitchHandlers (
   )\r
 {\r
   UINTN                           Index;\r
-  UINTN                           Bsp;\r
-  EXCEPTION_STACK_SWITCH_CONTEXT  SwitchStackData;\r
+  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;\r
   UINTN                           BufferSize;\r
+  EFI_STATUS                      Status;\r
+  UINT8                           *Buffer;\r
 \r
-  SwitchStackData.BufferSize = &BufferSize;\r
-  MpInitLibWhoAmI (&Bsp);\r
-\r
+  SwitchStackData = AllocateZeroPool (mNumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT));\r
+  ASSERT (SwitchStackData != NULL);\r
   for (Index = 0; Index < mNumberOfProcessors; ++Index) {\r
-    SwitchStackData.Buffer = NULL;\r
-    BufferSize             = 0;\r
+    //\r
+    // Because the procedure may runs multiple times, use the status EFI_NOT_STARTED\r
+    // to indicate the procedure haven't been run yet.\r
+    //\r
+    SwitchStackData[Index].Status = EFI_NOT_STARTED;\r
+  }\r
 \r
-    if (Index == Bsp) {\r
-      InitializeExceptionStackSwitchHandlers (&SwitchStackData);\r
+  Status = MpInitLibStartupAllCPUs (\r
+             InitializeExceptionStackSwitchHandlers,\r
+             0,\r
+             SwitchStackData\r
+             );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  BufferSize = 0;\r
+  for (Index = 0; Index < mNumberOfProcessors; ++Index) {\r
+    if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {\r
+      ASSERT (SwitchStackData[Index].BufferSize != 0);\r
+      BufferSize += SwitchStackData[Index].BufferSize;\r
     } else {\r
-      //\r
-      // AP might need different buffer size from BSP.\r
-      //\r
-      MpInitLibStartupThisAP (InitializeExceptionStackSwitchHandlers, Index, NULL, 0, (VOID *)&SwitchStackData, NULL);\r
+      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);\r
+      ASSERT (SwitchStackData[Index].BufferSize == 0);\r
     }\r
+  }\r
 \r
-    if (BufferSize == 0) {\r
-      continue;\r
+  if (BufferSize != 0) {\r
+    Buffer = AllocateRuntimeZeroPool (BufferSize);\r
+    ASSERT (Buffer != NULL);\r
+    BufferSize = 0;\r
+    for (Index = 0; Index < mNumberOfProcessors; ++Index) {\r
+      if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {\r
+        SwitchStackData[Index].Buffer = (VOID *)(&Buffer[BufferSize]);\r
+        BufferSize                   += SwitchStackData[Index].BufferSize;\r
+        DEBUG ((\r
+          DEBUG_INFO,\r
+          "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%lX\n",\r
+          (UINT64)(UINTN)Index,\r
+          (UINT64)(UINTN)SwitchStackData[Index].Buffer,\r
+          (UINT64)(UINTN)SwitchStackData[Index].BufferSize\r
+          ));\r
+      }\r
     }\r
 \r
-    SwitchStackData.Buffer = AllocateRuntimeZeroPool (BufferSize);\r
-    ASSERT (SwitchStackData.Buffer != NULL);\r
-    DEBUG ((\r
-      DEBUG_INFO,\r
-      "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%x\n",\r
-      (UINT64)(UINTN)Index,\r
-      (UINT64)(UINTN)SwitchStackData.Buffer,\r
-      (UINT32)BufferSize\r
-      ));\r
-\r
-    if (Index == Bsp) {\r
-      InitializeExceptionStackSwitchHandlers (&SwitchStackData);\r
-    } else {\r
-      MpInitLibStartupThisAP (\r
-        InitializeExceptionStackSwitchHandlers,\r
-        Index,\r
-        NULL,\r
-        0,\r
-        (VOID *)&SwitchStackData,\r
-        NULL\r
-        );\r
+    Status = MpInitLibStartupAllCPUs (\r
+               InitializeExceptionStackSwitchHandlers,\r
+               0,\r
+               SwitchStackData\r
+               );\r
+    ASSERT_EFI_ERROR (Status);\r
+    for (Index = 0; Index < mNumberOfProcessors; ++Index) {\r
+      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);\r
     }\r
   }\r
+\r
+  FreePool (SwitchStackData);\r
 }\r
 \r
 /**\r