]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
MdeModulePkg: Clean up source files
[mirror_edk2.git] / MdeModulePkg / Universal / Disk / DiskIoDxe / DiskIo.c
index e8af46f6c0ce60d9db7af0814d78fe3a243778ac..157e1b6d0ed95f424faaa10de83ecf28691bb8de 100644 (file)
@@ -9,7 +9,7 @@
     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
+Copyright (c) 2006 - 2018, 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
@@ -55,7 +55,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 +162,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 +171,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 +214,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 +289,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 +315,7 @@ DiskIoDriverBindingStop (
                     );\r
   }\r
   if (!EFI_ERROR (Status)) {\r
-    \r
+\r
     do {\r
       EfiAcquireLock (&Instance->TaskQueueLock);\r
       AllTaskDone = IsListEmpty (&Instance->TaskQueue);\r
@@ -324,7 +324,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 +343,7 @@ DiskIoDriverBindingStop (
                       );\r
       ASSERT_EFI_ERROR (Status);\r
     }\r
-    \r
+\r
     FreePool (Instance);\r
   }\r
 \r
@@ -366,11 +366,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 +411,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 +420,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 +438,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 +498,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 +553,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 +596,7 @@ DiskIoCreateSubtaskList (
       goto Done;\r
     }\r
     InsertTailList (Subtasks, &Subtask->Link);\r
-  \r
+\r
     BufferPtr  += Length;\r
     Offset     += Length;\r
     BufferSize -= Length;\r
@@ -642,13 +627,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 +655,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 +726,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 +752,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 +799,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 +817,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 +833,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 +854,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 +869,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 +890,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 +904,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 +913,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 +930,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 +956,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 +1088,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 +1138,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 +1235,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