[PATCH] UefiCpuPkg/PiSmmCpu: Don't allocate Token for SmmStartupThisAp

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

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

When Token points to mSmmStartupThisApToken, this routine is called
from SmmStartupThisAp() in non-blocking mode due to
PcdCpuSmmBlockStartupThisAp == FALSE.

In this case, caller wants to startup AP procedure in non-blocking
mode and cannot get the completion status from the Token because there
is no way to return the Token to caller from SmmStartupThisAp().
Caller needs to use its specific way to query the completion status.

There is no need to allocate a token for such case so the 2 overheads
can be avoided:
1. Call AllocateTokenBuffer() when there is no free token.
2. Get a free token from the token buffer.

Signed-off-by: Ray Ni <ray.ni@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 6227b2428a..efb89832ca 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -1,7 +1,7 @@
 /** @file
 SMM MP service implementation
 
-Copyright (c) 2009 - 2020, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2009 - 2021, Intel Corporation. All rights reserved.<BR>
 Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 
 SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -22,6 +22,7 @@ UINTN                                       mSemaphoreSize;
 SPIN_LOCK                                   *mPFLock = NULL;
 SMM_CPU_SYNC_MODE                           mCpuSmmSyncMode;
 BOOLEAN                                     mMachineCheckSupported = FALSE;
+MM_COMPLETION                               mSmmStartupThisApToken;
 
 extern UINTN mSmmShadowStackSize;
 
@@ -1240,9 +1241,23 @@ InternalSmmStartupThisAp (
   mSmmMpSyncData->CpuData[CpuIndex].Procedure = Procedure;
   mSmmMpSyncData->CpuData[CpuIndex].Parameter = ProcArguments;
   if (Token != NULL) {
-    ProcToken= GetFreeToken (1);
-    mSmmMpSyncData->CpuData[CpuIndex].Token = ProcToken;
-    *Token = (MM_COMPLETION)ProcToken->SpinLock;
+    if (Token != &mSmmStartupThisApToken) {
+      //
+      // When Token points to mSmmStartupThisApToken, this routine is called
+      // from SmmStartupThisAp() in non-blocking mode (PcdCpuSmmBlockStartupThisAp == FALSE).
+      // 
+      // In this case, caller wants to startup AP procedure in non-blocking
+      // mode and cannot get the completion status from the Token because there
+      // is no way to return the Token to caller from SmmStartupThisAp().
+      // Caller needs to use its implementation specific way to query the completion status.
+      // 
+      // There is no need to allocate a token for such case so the overhead of SMRAM and
+      // the allocation operation can be avoided.
+      //
+      ProcToken= GetFreeToken (1);
+      mSmmMpSyncData->CpuData[CpuIndex].Token = ProcToken;
+      *Token = (MM_COMPLETION)ProcToken->SpinLock;
+    }
   }
   mSmmMpSyncData->CpuData[CpuIndex].Status    = CpuStatus;
   if (mSmmMpSyncData->CpuData[CpuIndex].Status != NULL) {
@@ -1474,8 +1489,6 @@ SmmStartupThisAp (
   IN OUT  VOID                      *ProcArguments OPTIONAL
   )
 {
-  MM_COMPLETION               Token;
-
   gSmmCpuPrivate->ApWrapperFunc[CpuIndex].Procedure = Procedure;
   gSmmCpuPrivate->ApWrapperFunc[CpuIndex].ProcedureArgument = ProcArguments;
 
@@ -1486,7 +1499,7 @@ SmmStartupThisAp (
     ProcedureWrapper,
     CpuIndex,
     &gSmmCpuPrivate->ApWrapperFunc[CpuIndex],
-    FeaturePcdGet (PcdCpuSmmBlockStartupThisAp) ? NULL : &Token,
+    FeaturePcdGet (PcdCpuSmmBlockStartupThisAp) ? NULL : &mSmmStartupThisApToken,
     0,
     NULL
     );
-- 
2.27.0.windows.1