]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
MdeModulePkg: Replace BSD License with BSD+Patent License
[mirror_edk2.git] / MdeModulePkg / Universal / Disk / DiskIoDxe / DiskIo.c
index e8af46f6c0ce60d9db7af0814d78fe3a243778ac..e19466bd26635c4285773b2f2e0cbc905ed53070 100644 (file)
@@ -9,14 +9,8 @@
     Aligned  - A read of N contiguous sectors.\r
     OverRun  - The last byte is not on a sector boundary.\r
 \r
-Copyright (c) 2006 - 2013, 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
-http://opensource.org/licenses/bsd-license.php\r
-\r
-THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+Copyright (c) 2006 - 2018, Intel Corporation. All rights reserved.<BR>\r
+SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
@@ -55,7 +49,7 @@ DISK_IO_PRIVATE_DATA        gDiskIoPrivateDataTemplate = {
 };\r
 \r
 /**\r
-  Test to see if this driver supports ControllerHandle. \r
+  Test to see if this driver supports ControllerHandle.\r
 \r
   @param  This                Protocol instance pointer.\r
   @param  ControllerHandle    Handle of device to test\r
@@ -162,7 +156,7 @@ DiskIoDriverBindingStart (
   if (EFI_ERROR (Status)) {\r
     gDiskIoPrivateDataTemplate.BlockIo2 = NULL;\r
   }\r
-  \r
+\r
   //\r
   // Initialize the Disk IO device instance.\r
   //\r
@@ -171,19 +165,19 @@ DiskIoDriverBindingStart (
     Status = EFI_OUT_OF_RESOURCES;\r
     goto ErrorExit;\r
   }\r
-  \r
+\r
   //\r
   // The BlockSize and IoAlign of BlockIo and BlockIo2 should equal.\r
   //\r
   ASSERT ((Instance->BlockIo2 == NULL) ||\r
-          ((Instance->BlockIo->Media->IoAlign == Instance->BlockIo2->Media->IoAlign) && \r
+          ((Instance->BlockIo->Media->IoAlign == Instance->BlockIo2->Media->IoAlign) &&\r
            (Instance->BlockIo->Media->BlockSize == Instance->BlockIo2->Media->BlockSize)\r
           ));\r
-  \r
+\r
   InitializeListHead (&Instance->TaskQueue);\r
   EfiInitializeLock (&Instance->TaskQueueLock, TPL_NOTIFY);\r
   Instance->SharedWorkingBuffer = AllocateAlignedPages (\r
-                                    EFI_SIZE_TO_PAGES (DATA_BUFFER_BLOCK_NUM * Instance->BlockIo->Media->BlockSize),\r
+                                    EFI_SIZE_TO_PAGES (PcdGet32 (PcdDiskIoDataBufferBlockNum) * Instance->BlockIo->Media->BlockSize),\r
                                     Instance->BlockIo->Media->IoAlign\r
                                     );\r
   if (Instance->SharedWorkingBuffer == NULL) {\r
@@ -214,7 +208,7 @@ ErrorExit:
     if (Instance != NULL && Instance->SharedWorkingBuffer != NULL) {\r
       FreeAlignedPages (\r
         Instance->SharedWorkingBuffer,\r
-        EFI_SIZE_TO_PAGES (DATA_BUFFER_BLOCK_NUM * Instance->BlockIo->Media->BlockSize)\r
+        EFI_SIZE_TO_PAGES (PcdGet32 (PcdDiskIoDataBufferBlockNum) * Instance->BlockIo->Media->BlockSize)\r
         );\r
     }\r
 \r
@@ -289,7 +283,7 @@ DiskIoDriverBindingStop (
   if (EFI_ERROR (Status)) {\r
     DiskIo2 = NULL;\r
   }\r
-  \r
+\r
   Instance = DISK_IO_PRIVATE_DATA_FROM_DISK_IO (DiskIo);\r
 \r
   if (DiskIo2 != NULL) {\r
@@ -315,7 +309,7 @@ DiskIoDriverBindingStop (
                     );\r
   }\r
   if (!EFI_ERROR (Status)) {\r
-    \r
+\r
     do {\r
       EfiAcquireLock (&Instance->TaskQueueLock);\r
       AllTaskDone = IsListEmpty (&Instance->TaskQueue);\r
@@ -324,7 +318,7 @@ DiskIoDriverBindingStop (
 \r
     FreeAlignedPages (\r
       Instance->SharedWorkingBuffer,\r
-      EFI_SIZE_TO_PAGES (DATA_BUFFER_BLOCK_NUM * Instance->BlockIo->Media->BlockSize)\r
+      EFI_SIZE_TO_PAGES (PcdGet32 (PcdDiskIoDataBufferBlockNum) * Instance->BlockIo->Media->BlockSize)\r
       );\r
 \r
     Status = gBS->CloseProtocol (\r
@@ -343,7 +337,7 @@ DiskIoDriverBindingStop (
                       );\r
       ASSERT_EFI_ERROR (Status);\r
     }\r
-    \r
+\r
     FreePool (Instance);\r
   }\r
 \r
@@ -366,11 +360,19 @@ DiskIoDestroySubtask (
   )\r
 {\r
   LIST_ENTRY               *Link;\r
+\r
+  if (Subtask->Task != NULL) {\r
+    EfiAcquireLock (&Subtask->Task->SubtasksLock);\r
+  }\r
   Link = RemoveEntryList (&Subtask->Link);\r
+  if (Subtask->Task != NULL) {\r
+    EfiReleaseLock (&Subtask->Task->SubtasksLock);\r
+  }\r
+\r
   if (!Subtask->Blocking) {\r
     if (Subtask->WorkingBuffer != NULL) {\r
       FreeAlignedPages (\r
-        Subtask->WorkingBuffer, \r
+        Subtask->WorkingBuffer,\r
         Subtask->Length < Instance->BlockIo->Media->BlockSize\r
         ? EFI_SIZE_TO_PAGES (Instance->BlockIo->Media->BlockSize)\r
         : EFI_SIZE_TO_PAGES (Subtask->Length)\r
@@ -403,8 +405,6 @@ DiskIo2OnReadWriteComplete (
   EFI_STATUS            TransactionStatus;\r
   DISK_IO_PRIVATE_DATA  *Instance;\r
 \r
-  gBS->CloseEvent (Event);\r
-\r
   Subtask           = (DISK_IO_SUBTASK *) Context;\r
   TransactionStatus = Subtask->BlockIo2Token.TransactionStatus;\r
   Task              = Subtask->Task;\r
@@ -414,26 +414,13 @@ DiskIo2OnReadWriteComplete (
   ASSERT (Instance->Signature == DISK_IO_PRIVATE_DATA_SIGNATURE);\r
   ASSERT (Task->Signature     == DISK_IO2_TASK_SIGNATURE);\r
 \r
-  if (Subtask->WorkingBuffer != NULL) {\r
-    if (!EFI_ERROR (TransactionStatus) && (Task->Token != NULL) && !Subtask->Write) {\r
-      CopyMem (Subtask->Buffer, Subtask->WorkingBuffer + Subtask->Offset, Subtask->Length);\r
-    }\r
-\r
-    //\r
-    // The WorkingBuffer of blocking subtask either points to SharedWorkingBuffer\r
-    // or will be used by non-blocking subtask which will be freed below.\r
-    //\r
-    if (!Subtask->Blocking) {\r
-      FreeAlignedPages (\r
-        Subtask->WorkingBuffer, \r
-        Subtask->Length < Instance->BlockIo->Media->BlockSize\r
-        ? EFI_SIZE_TO_PAGES (Instance->BlockIo->Media->BlockSize)\r
-        : EFI_SIZE_TO_PAGES (Subtask->Length)\r
-        );\r
-    }\r
+  if ((Subtask->WorkingBuffer != NULL) && !EFI_ERROR (TransactionStatus) &&\r
+      (Task->Token != NULL) && !Subtask->Write\r
+     ) {\r
+    CopyMem (Subtask->Buffer, Subtask->WorkingBuffer + Subtask->Offset, Subtask->Length);\r
   }\r
-  RemoveEntryList (&Subtask->Link);\r
-  FreePool (Subtask);\r
+\r
+  DiskIoDestroySubtask (Instance, Subtask);\r
 \r
   if (EFI_ERROR (TransactionStatus) || IsListEmpty (&Task->Subtasks)) {\r
     if (Task->Token != NULL) {\r
@@ -445,19 +432,11 @@ DiskIo2OnReadWriteComplete (
       gBS->SignalEvent (Task->Token->Event);\r
 \r
       //\r
-      // Mark token to NULL\r
+      // Mark token to NULL indicating the Task is a dead task.\r
       //\r
       Task->Token = NULL;\r
     }\r
   }\r
-\r
-  if (IsListEmpty (&Task->Subtasks)) {\r
-    EfiAcquireLock (&Instance->TaskQueueLock);\r
-    RemoveEntryList (&Task->Link);\r
-    EfiReleaseLock (&Instance->TaskQueueLock);\r
-\r
-    FreePool (Task);\r
-  }\r
 }\r
 \r
 /**\r
@@ -513,7 +492,7 @@ DiskIoCreateSubtask (
     }\r
   }\r
   DEBUG ((\r
-    EFI_D_BLKIO, \r
+    EFI_D_BLKIO,\r
     "  %c:Lba/Offset/Length/WorkingBuffer/Buffer = %016lx/%08x/%08x/%08x/%08x\n",\r
     Write ? 'W': 'R', Lba, Offset, Length, WorkingBuffer, Buffer\r
     ));\r
@@ -568,7 +547,7 @@ DiskIoCreateSubtaskList (
   if (IoAlign == 0) {\r
     IoAlign = 1;\r
   }\r
-  \r
+\r
   Lba       = DivU64x32Remainder (Offset, BlockSize, &UnderRun);\r
   BufferPtr = (UINT8 *) Buffer;\r
 \r
@@ -611,7 +590,7 @@ DiskIoCreateSubtaskList (
       goto Done;\r
     }\r
     InsertTailList (Subtasks, &Subtask->Link);\r
-  \r
+\r
     BufferPtr  += Length;\r
     Offset     += Length;\r
     BufferSize -= Length;\r
@@ -642,13 +621,13 @@ DiskIoCreateSubtaskList (
       InsertTailList (Subtasks, &Subtask->Link);\r
     }\r
 \r
-    Subtask = DiskIoCreateSubtask (Write, OverRunLba, 0, OverRun, WorkingBuffer, BufferPtr, Blocking);\r
+    Subtask = DiskIoCreateSubtask (Write, OverRunLba, 0, OverRun, WorkingBuffer, BufferPtr + BufferSize, Blocking);\r
     if (Subtask == NULL) {\r
       goto Done;\r
     }\r
     InsertTailList (Subtasks, &Subtask->Link);\r
   }\r
-  \r
+\r
   if (OverRunLba > Lba) {\r
     //\r
     // If the DiskIo maps directly to a BlockIo device do the read.\r
@@ -670,8 +649,8 @@ DiskIoCreateSubtaskList (
         // Use the allocated buffer instead of the original buffer\r
         // to avoid alignment issue.\r
         //\r
-        for (; Lba < OverRunLba; Lba += DATA_BUFFER_BLOCK_NUM) {\r
-          DataBufferSize = MIN (BufferSize, DATA_BUFFER_BLOCK_NUM * BlockSize);\r
+        for (; Lba < OverRunLba; Lba += PcdGet32 (PcdDiskIoDataBufferBlockNum)) {\r
+          DataBufferSize = MIN (BufferSize, PcdGet32 (PcdDiskIoDataBufferBlockNum) * BlockSize);\r
 \r
           Subtask = DiskIoCreateSubtask (Write, Lba, 0, DataBufferSize, SharedWorkingBuffer, BufferPtr, Blocking);\r
           if (Subtask == NULL) {\r
@@ -741,7 +720,7 @@ DiskIo2Cancel (
   DISK_IO_PRIVATE_DATA  *Instance;\r
   DISK_IO2_TASK         *Task;\r
   LIST_ENTRY            *Link;\r
-  \r
+\r
   Instance = DISK_IO_PRIVATE_DATA_FROM_DISK_IO2 (This);\r
 \r
   EfiAcquireLock (&Instance->TaskQueueLock);\r
@@ -767,6 +746,42 @@ DiskIo2Cancel (
   return EFI_SUCCESS;\r
 }\r
 \r
+/**\r
+  Remove the completed tasks from Instance->TaskQueue. Completed tasks are those who don't have any subtasks.\r
+\r
+  @param Instance    Pointer to the DISK_IO_PRIVATE_DATA.\r
+\r
+  @retval TRUE       The Instance->TaskQueue is empty after the completed tasks are removed.\r
+  @retval FALSE      The Instance->TaskQueue is not empty after the completed tasks are removed.\r
+**/\r
+BOOLEAN\r
+DiskIo2RemoveCompletedTask (\r
+  IN DISK_IO_PRIVATE_DATA     *Instance\r
+  )\r
+{\r
+  BOOLEAN                     QueueEmpty;\r
+  LIST_ENTRY                  *Link;\r
+  DISK_IO2_TASK               *Task;\r
+\r
+  QueueEmpty = TRUE;\r
+\r
+  EfiAcquireLock (&Instance->TaskQueueLock);\r
+  for (Link = GetFirstNode (&Instance->TaskQueue); !IsNull (&Instance->TaskQueue, Link); ) {\r
+    Task = CR (Link, DISK_IO2_TASK, Link, DISK_IO2_TASK_SIGNATURE);\r
+    if (IsListEmpty (&Task->Subtasks)) {\r
+      Link = RemoveEntryList (&Task->Link);\r
+      ASSERT (Task->Token == NULL);\r
+      FreePool (Task);\r
+    } else {\r
+      Link = GetNextNode (&Instance->TaskQueue, Link);\r
+      QueueEmpty = FALSE;\r
+    }\r
+  }\r
+  EfiReleaseLock (&Instance->TaskQueueLock);\r
+\r
+  return QueueEmpty;\r
+}\r
+\r
 /**\r
   Common routine to access the disk.\r
 \r
@@ -778,7 +793,7 @@ DiskIo2Cancel (
                      If this field is NULL, synchronous/blocking IO is performed.\r
   @param  BufferSize            The size in bytes of Buffer. The number of bytes to read from the device.\r
   @param  Buffer                A pointer to the destination buffer for the data.\r
-                                The caller is responsible either having implicit or explicit ownership of the buffer. \r
+                                The caller is responsible either having implicit or explicit ownership of the buffer.\r
 **/\r
 EFI_STATUS\r
 DiskIo2ReadWriteDisk (\r
@@ -796,12 +811,13 @@ DiskIo2ReadWriteDisk (
   EFI_BLOCK_IO2_PROTOCOL *BlockIo2;\r
   EFI_BLOCK_IO_MEDIA     *Media;\r
   LIST_ENTRY             *Link;\r
+  LIST_ENTRY             *NextLink;\r
   LIST_ENTRY             Subtasks;\r
   DISK_IO_SUBTASK        *Subtask;\r
   DISK_IO2_TASK          *Task;\r
-  BOOLEAN                TaskQueueEmpty;\r
   EFI_TPL                OldTpl;\r
   BOOLEAN                Blocking;\r
+  BOOLEAN                SubtaskBlocking;\r
   LIST_ENTRY             *SubtasksPtr;\r
 \r
   Task      = NULL;\r
@@ -811,31 +827,20 @@ DiskIo2ReadWriteDisk (
   Status    = EFI_SUCCESS;\r
   Blocking  = (BOOLEAN) ((Token == NULL) || (Token->Event == NULL));\r
 \r
-  if (Media->MediaId != MediaId) {\r
-    return EFI_MEDIA_CHANGED;\r
-  }\r
-  \r
-  if (Write && Media->ReadOnly) {\r
-    return EFI_WRITE_PROTECTED;\r
-  }\r
-  \r
   if (Blocking) {\r
     //\r
     // Wait till pending async task is completed.\r
     //\r
-    do {\r
-      EfiAcquireLock (&Instance->TaskQueueLock);\r
-      TaskQueueEmpty = IsListEmpty (&Instance->TaskQueue);\r
-      EfiReleaseLock (&Instance->TaskQueueLock);\r
-    } while (!TaskQueueEmpty);\r
+    while (!DiskIo2RemoveCompletedTask (Instance));\r
 \r
     SubtasksPtr = &Subtasks;\r
   } else {\r
+    DiskIo2RemoveCompletedTask (Instance);\r
     Task = AllocatePool (sizeof (DISK_IO2_TASK));\r
     if (Task == NULL) {\r
       return EFI_OUT_OF_RESOURCES;\r
     }\r
-    \r
+\r
     EfiAcquireLock (&Instance->TaskQueueLock);\r
     InsertTailList (&Instance->TaskQueue, &Task->Link);\r
     EfiReleaseLock (&Instance->TaskQueueLock);\r
@@ -843,6 +848,7 @@ DiskIo2ReadWriteDisk (
     Task->Signature = DISK_IO2_TASK_SIGNATURE;\r
     Task->Instance  = Instance;\r
     Task->Token     = Token;\r
+    EfiInitializeLock (&Task->SubtasksLock, TPL_NOTIFY);\r
 \r
     SubtasksPtr = &Task->Subtasks;\r
   }\r
@@ -857,12 +863,15 @@ DiskIo2ReadWriteDisk (
   ASSERT (!IsListEmpty (SubtasksPtr));\r
 \r
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
-  for (Link = GetFirstNode (SubtasksPtr); !IsNull (SubtasksPtr, Link); Link = GetNextNode (SubtasksPtr, Link)) {\r
-    Subtask = CR (Link, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
+  for ( Link = GetFirstNode (SubtasksPtr), NextLink = GetNextNode (SubtasksPtr, Link)\r
+      ; !IsNull (SubtasksPtr, Link)\r
+      ; Link = NextLink, NextLink = GetNextNode (SubtasksPtr, NextLink)\r
+      ) {\r
+    Subtask         = CR (Link, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
+    Subtask->Task   = Task;\r
+    SubtaskBlocking = Subtask->Blocking;\r
 \r
-    if (!Subtask->Blocking) {\r
-      Subtask->Task = Task;\r
-    }\r
+    ASSERT ((Subtask->Length % Media->BlockSize == 0) || (Subtask->Length < Media->BlockSize));\r
 \r
     if (Subtask->Write) {\r
       //\r
@@ -875,12 +884,12 @@ DiskIo2ReadWriteDisk (
         CopyMem (Subtask->WorkingBuffer + Subtask->Offset, Subtask->Buffer, Subtask->Length);\r
       }\r
 \r
-      if (Subtask->Blocking) {\r
+      if (SubtaskBlocking) {\r
         Status = BlockIo->WriteBlocks (\r
                             BlockIo,\r
                             MediaId,\r
                             Subtask->Lba,\r
-                            Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                            (Subtask->Length % Media->BlockSize == 0) ? Subtask->Length : Media->BlockSize,\r
                             (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
                             );\r
       } else {\r
@@ -889,7 +898,7 @@ DiskIo2ReadWriteDisk (
                              MediaId,\r
                              Subtask->Lba,\r
                              &Subtask->BlockIo2Token,\r
-                             Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                             (Subtask->Length % Media->BlockSize == 0) ? Subtask->Length : Media->BlockSize,\r
                              (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
                              );\r
       }\r
@@ -898,12 +907,12 @@ DiskIo2ReadWriteDisk (
       //\r
       // Read\r
       //\r
-      if (Subtask->Blocking) {\r
+      if (SubtaskBlocking) {\r
         Status = BlockIo->ReadBlocks (\r
                             BlockIo,\r
                             MediaId,\r
                             Subtask->Lba,\r
-                            Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                            (Subtask->Length % Media->BlockSize == 0) ? Subtask->Length : Media->BlockSize,\r
                             (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
                             );\r
         if (!EFI_ERROR (Status) && (Subtask->WorkingBuffer != NULL)) {\r
@@ -915,17 +924,25 @@ DiskIo2ReadWriteDisk (
                              MediaId,\r
                              Subtask->Lba,\r
                              &Subtask->BlockIo2Token,\r
-                             Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                             (Subtask->Length % Media->BlockSize == 0) ? Subtask->Length : Media->BlockSize,\r
                              (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
                              );\r
       }\r
     }\r
-    \r
+\r
+    if (SubtaskBlocking || EFI_ERROR (Status)) {\r
+      //\r
+      // Make sure the subtask list only contains non-blocking subtasks.\r
+      // Remove failed non-blocking subtasks as well because the callback won't be called.\r
+      //\r
+      DiskIoDestroySubtask (Instance, Subtask);\r
+    }\r
+\r
     if (EFI_ERROR (Status)) {\r
       break;\r
     }\r
   }\r
-  \r
+\r
   gBS->RaiseTPL (TPL_NOTIFY);\r
 \r
   //\r
@@ -933,35 +950,22 @@ DiskIo2ReadWriteDisk (
   // We shouldn't remove all the tasks because the non-blocking requests have been submitted and cannot be canceled.\r
   //\r
   if (EFI_ERROR (Status)) {\r
-    while (!IsNull (SubtasksPtr, Link)) {\r
-      Subtask = CR (Link, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
-      Link = DiskIoDestroySubtask (Instance, Subtask);\r
-    }\r
-  }\r
-\r
-  //\r
-  // Remove all the blocking subtasks because the non-blocking callback only removes the non-blocking subtask.\r
-  //\r
-  for (Link = GetFirstNode (SubtasksPtr); !IsNull (SubtasksPtr, Link); ) {\r
-    Subtask = CR (Link, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
-    if (Subtask->Blocking) {\r
-      Link = DiskIoDestroySubtask (Instance, Subtask);\r
-    } else {\r
-      Link = GetNextNode (SubtasksPtr, Link);\r
+    while (!IsNull (SubtasksPtr, NextLink)) {\r
+      Subtask = CR (NextLink, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
+      NextLink = DiskIoDestroySubtask (Instance, Subtask);\r
     }\r
   }\r
 \r
   //\r
-  // It's possible that the callback runs before raising TPL to NOTIFY,\r
-  // so the subtasks list only contains blocking subtask.\r
-  // Remove the Task after the blocking subtasks are removed in above.\r
+  // It's possible that the non-blocking subtasks finish before raising TPL to NOTIFY,\r
+  // so the subtasks list might be empty at this point.\r
   //\r
   if (!Blocking && IsListEmpty (SubtasksPtr)) {\r
     EfiAcquireLock (&Instance->TaskQueueLock);\r
     RemoveEntryList (&Task->Link);\r
     EfiReleaseLock (&Instance->TaskQueueLock);\r
 \r
-    if (Task->Token != NULL) {\r
+    if (!EFI_ERROR (Status) && (Task->Token != NULL)) {\r
       //\r
       // Task->Token should be set to NULL by the DiskIo2OnReadWriteComplete\r
       // It it's not, that means the non-blocking request was downgraded to blocking request.\r
@@ -1078,6 +1082,8 @@ DiskIo2OnFlushComplete (
   ASSERT (Task->Signature == DISK_IO2_FLUSH_TASK_SIGNATURE);\r
   Task->Token->TransactionStatus = Task->BlockIo2Token.TransactionStatus;\r
   gBS->SignalEvent (Task->Token->Event);\r
+\r
+  FreePool (Task);\r
 }\r
 \r
 /**\r
@@ -1126,6 +1132,7 @@ DiskIo2FlushDiskEx (
       return Status;\r
     }\r
     Task->Signature = DISK_IO2_FLUSH_TASK_SIGNATURE;\r
+    Task->Token     = Token;\r
     Status = Private->BlockIo2->FlushBlocksEx (Private->BlockIo2, &Task->BlockIo2Token);\r
     if (EFI_ERROR (Status)) {\r
       gBS->CloseEvent (Task->BlockIo2Token.Event);\r
@@ -1222,9 +1229,9 @@ DiskIoWriteDisk (
 /**\r
   The user Entry Point for module DiskIo. The user code starts with this function.\r
 \r
-  @param[in] ImageHandle    The firmware allocated handle for the EFI image.  \r
+  @param[in] ImageHandle    The firmware allocated handle for the EFI image.\r
   @param[in] SystemTable    A pointer to the EFI System Table.\r
-  \r
+\r
   @retval EFI_SUCCESS       The entry point is executed successfully.\r
   @retval other             Some error occurs when executing this entry point.\r
 \r