]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
Add DiskIo2 protocol definition to MdePkg.
[mirror_edk2.git] / MdeModulePkg / Universal / Disk / DiskIoDxe / DiskIo.c
index 0cb49224ca7ad80c13ad43cc0340482a05378976..4bd6772046310cd2e367e11dfa321012d49ea9fb 100644 (file)
@@ -1,27 +1,30 @@
 /** @file\r
-  DiskIo driver that layers it's self on every Block IO protocol in the system.\r
+  DiskIo driver that lays on every BlockIo protocol in the system.\r
   DiskIo converts a block oriented device to a byte oriented device.\r
 \r
-  ReadDisk may have to do reads that are not aligned on sector boundaries.\r
+  Disk access may have to handle unaligned request about sector boundaries.\r
   There are three cases:\r
     UnderRun - The first byte is not on a sector boundary or the read request is\r
                less than a sector in length.\r
     Aligned  - A read of N contiguous sectors.\r
     OverRun  - The last byte is not on a sector boundary.\r
 \r
-  Copyright (c) 2006 - 2007, Intel Corporation                                              \r
-  All rights reserved. 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
+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
+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
 \r
 **/\r
 \r
 #include "DiskIo.h"\r
 \r
+//\r
+// Driver binding protocol implementation for DiskIo driver.\r
+//\r
 EFI_DRIVER_BINDING_PROTOCOL gDiskIoDriverBinding = {\r
   DiskIoDriverBindingSupported,\r
   DiskIoDriverBindingStart,\r
@@ -31,6 +34,10 @@ EFI_DRIVER_BINDING_PROTOCOL gDiskIoDriverBinding = {
   NULL\r
 };\r
 \r
+//\r
+// Template for DiskIo private data structure.\r
+// The pointer to BlockIo protocol interface is assigned dynamically.\r
+//\r
 DISK_IO_PRIVATE_DATA        gDiskIoPrivateDataTemplate = {\r
   DISK_IO_PRIVATE_DATA_SIGNATURE,\r
   {\r
@@ -38,10 +45,15 @@ DISK_IO_PRIVATE_DATA        gDiskIoPrivateDataTemplate = {
     DiskIoReadDisk,\r
     DiskIoWriteDisk\r
   },\r
-  NULL\r
+  {\r
+    EFI_DISK_IO2_PROTOCOL_REVISION,\r
+    DiskIo2Cancel,\r
+    DiskIo2ReadDiskEx,\r
+    DiskIo2WriteDiskEx,\r
+    DiskIo2FlushDiskEx\r
+  }\r
 };\r
 \r
-\r
 /**\r
   Test to see if this driver supports ControllerHandle. \r
 \r
@@ -85,11 +97,11 @@ DiskIoDriverBindingSupported (
   // Close the I/O Abstraction(s) used to perform the supported test.\r
   //\r
   gBS->CloseProtocol (\r
-        ControllerHandle,\r
-        &gEfiBlockIoProtocolGuid,\r
-        This->DriverBindingHandle,\r
-        ControllerHandle\r
-        );\r
+         ControllerHandle,\r
+         &gEfiBlockIoProtocolGuid,\r
+         This->DriverBindingHandle,\r
+         ControllerHandle\r
+         );\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -117,12 +129,16 @@ DiskIoDriverBindingStart (
   )\r
 {\r
   EFI_STATUS            Status;\r
-  DISK_IO_PRIVATE_DATA  *Private;\r
+  DISK_IO_PRIVATE_DATA  *Instance;\r
+  EFI_TPL               OldTpl;\r
+  EFI_BLOCK_IO_MEDIA    *Media;\r
 \r
-  Private = NULL;\r
+  Instance = NULL;\r
+\r
+  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
 \r
   //\r
-  // Connect to the Block IO interface on ControllerHandle.\r
+  // Connect to the Block IO and Block IO2 interface on ControllerHandle.\r
   //\r
   Status = gBS->OpenProtocol (\r
                   ControllerHandle,\r
@@ -133,47 +149,95 @@ DiskIoDriverBindingStart (
                   EFI_OPEN_PROTOCOL_BY_DRIVER\r
                   );\r
   if (EFI_ERROR (Status)) {\r
-    return Status;\r
+    goto ErrorExit1;\r
+  }\r
+\r
+  Status = gBS->OpenProtocol (\r
+                  ControllerHandle,\r
+                  &gEfiBlockIo2ProtocolGuid,\r
+                  (VOID **) &gDiskIoPrivateDataTemplate.BlockIo2,\r
+                  This->DriverBindingHandle,\r
+                  ControllerHandle,\r
+                  EFI_OPEN_PROTOCOL_BY_DRIVER\r
+                  );\r
+  if (EFI_ERROR (Status)) {\r
+    gDiskIoPrivateDataTemplate.BlockIo2 = NULL;\r
   }\r
   \r
   //\r
   // Initialize the Disk IO device instance.\r
   //\r
-  Private = AllocateCopyPool (sizeof (DISK_IO_PRIVATE_DATA), &gDiskIoPrivateDataTemplate);\r
-  if (Private == NULL) {\r
+  Instance = AllocateCopyPool (sizeof (DISK_IO_PRIVATE_DATA), &gDiskIoPrivateDataTemplate);\r
+  if (Instance == NULL) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
     goto ErrorExit;\r
   }\r
   \r
+  Media = Instance->BlockIo->Media;\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->BlockSize == Instance->BlockIo2->Media->BlockSize)\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
+                                    Instance->BlockIo->Media->IoAlign\r
+                                    );\r
+  if (Instance->SharedWorkingBuffer == NULL) {\r
+    Status = EFI_OUT_OF_RESOURCES;\r
+    goto ErrorExit;\r
+  }\r
+\r
   //\r
   // Install protocol interfaces for the Disk IO device.\r
   //\r
-  Status = gBS->InstallProtocolInterface (\r
-                  &ControllerHandle,\r
-                  &gEfiDiskIoProtocolGuid,\r
-                  EFI_NATIVE_INTERFACE,\r
-                  &Private->DiskIo\r
-                  );\r
+  if (Instance->BlockIo2 != NULL) {\r
+    Status = gBS->InstallMultipleProtocolInterfaces (\r
+                    &ControllerHandle,\r
+                    &gEfiDiskIoProtocolGuid,  &Instance->DiskIo,\r
+                    &gEfiDiskIo2ProtocolGuid, &Instance->DiskIo2,\r
+                    NULL\r
+                    );\r
+  } else {\r
+    Status = gBS->InstallMultipleProtocolInterfaces (\r
+                    &ControllerHandle,\r
+                    &gEfiDiskIoProtocolGuid,  &Instance->DiskIo,\r
+                    NULL\r
+                    );\r
+  }\r
 \r
 ErrorExit:\r
   if (EFI_ERROR (Status)) {\r
+    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
+        );\r
+    }\r
 \r
-    if (Private != NULL) {\r
-      FreePool (Private);\r
+    if (Instance != NULL) {\r
+      FreePool (Instance);\r
     }\r
 \r
     gBS->CloseProtocol (\r
-          ControllerHandle,\r
-          &gEfiBlockIoProtocolGuid,\r
-          This->DriverBindingHandle,\r
-          ControllerHandle\r
-          );\r
+           ControllerHandle,\r
+           &gEfiBlockIoProtocolGuid,\r
+           This->DriverBindingHandle,\r
+           ControllerHandle\r
+           );\r
   }\r
 \r
+ErrorExit1:\r
+  gBS->RestoreTPL (OldTpl);\r
   return Status;\r
 }\r
 \r
-\r
 /**\r
   Stop this driver on ControllerHandle by removing Disk IO protocol and closing\r
   the Block IO protocol on ControllerHandle.\r
@@ -199,7 +263,9 @@ DiskIoDriverBindingStop (
 {\r
   EFI_STATUS            Status;\r
   EFI_DISK_IO_PROTOCOL  *DiskIo;\r
-  DISK_IO_PRIVATE_DATA  *Private;\r
+  EFI_DISK_IO2_PROTOCOL *DiskIo2;\r
+  DISK_IO_PRIVATE_DATA  *Instance;\r
+  BOOLEAN               AllTaskDone;\r
 \r
   //\r
   // Get our context back.\r
@@ -213,17 +279,56 @@ DiskIoDriverBindingStop (
                   EFI_OPEN_PROTOCOL_GET_PROTOCOL\r
                   );\r
   if (EFI_ERROR (Status)) {\r
-    return EFI_UNSUPPORTED;\r
+    return Status;\r
   }\r
-\r
-  Private = DISK_IO_PRIVATE_DATA_FROM_THIS (DiskIo);\r
-\r
-  Status = gBS->UninstallProtocolInterface (\r
+  Status = gBS->OpenProtocol (\r
                   ControllerHandle,\r
-                  &gEfiDiskIoProtocolGuid,\r
-                  &Private->DiskIo\r
+                  &gEfiDiskIo2ProtocolGuid,\r
+                  (VOID **) &DiskIo2,\r
+                  This->DriverBindingHandle,\r
+                  ControllerHandle,\r
+                  EFI_OPEN_PROTOCOL_GET_PROTOCOL\r
                   );\r
+  if (EFI_ERROR (Status)) {\r
+    DiskIo2 = NULL;\r
+  }\r
+  \r
+  Instance = DISK_IO_PRIVATE_DATA_FROM_DISK_IO (DiskIo);\r
+\r
+  if (DiskIo2 != NULL) {\r
+    //\r
+    // Call BlockIo2::Reset() to terminate any in-flight non-blocking I/O requests\r
+    //\r
+    ASSERT (Instance->BlockIo2 != NULL);\r
+    Status = Instance->BlockIo2->Reset (Instance->BlockIo2, FALSE);\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+    Status = gBS->UninstallMultipleProtocolInterfaces (\r
+                    ControllerHandle,\r
+                    &gEfiDiskIoProtocolGuid,  &Instance->DiskIo,\r
+                    &gEfiDiskIo2ProtocolGuid, &Instance->DiskIo2,\r
+                    NULL\r
+                    );\r
+  } else {\r
+    Status = gBS->UninstallMultipleProtocolInterfaces (\r
+                    ControllerHandle,\r
+                    &gEfiDiskIoProtocolGuid,  &Instance->DiskIo,\r
+                    NULL\r
+                    );\r
+  }\r
   if (!EFI_ERROR (Status)) {\r
+    \r
+    do {\r
+      EfiAcquireLock (&Instance->TaskQueueLock);\r
+      AllTaskDone = IsListEmpty (&Instance->TaskQueue);\r
+      EfiReleaseLock (&Instance->TaskQueueLock);\r
+    } while (!AllTaskDone);\r
+\r
+    FreeAlignedPages (\r
+      Instance->SharedWorkingBuffer,\r
+      EFI_SIZE_TO_PAGES (DATA_BUFFER_BLOCK_NUM * Instance->BlockIo->Media->BlockSize)\r
+      );\r
 \r
     Status = gBS->CloseProtocol (\r
                     ControllerHandle,\r
@@ -231,472 +336,890 @@ DiskIoDriverBindingStop (
                     This->DriverBindingHandle,\r
                     ControllerHandle\r
                     );\r
+    ASSERT_EFI_ERROR (Status);\r
+    if (DiskIo2 != NULL) {\r
+      Status = gBS->CloseProtocol (\r
+                      ControllerHandle,\r
+                      &gEfiBlockIo2ProtocolGuid,\r
+                      This->DriverBindingHandle,\r
+                      ControllerHandle\r
+                      );\r
+      ASSERT_EFI_ERROR (Status);\r
+    }\r
+    \r
+    FreePool (Instance);\r
   }\r
 \r
-  if (!EFI_ERROR (Status)) {\r
-    FreePool (Private);\r
+  return Status;\r
+}\r
+\r
+\r
+/**\r
+  Destroy the sub task.\r
+\r
+  @param Subtask      Subtask.\r
+\r
+  @return LIST_ENTRY *  Pointer to the next link of subtask.\r
+**/\r
+LIST_ENTRY *\r
+DiskIoDestroySubtask (\r
+  IN DISK_IO_PRIVATE_DATA     *Instance,\r
+  IN DISK_IO_SUBTASK          *Subtask\r
+  )\r
+{\r
+  LIST_ENTRY               *Link;\r
+  Link = RemoveEntryList (&Subtask->Link);\r
+  if (!Subtask->Blocking) {\r
+    if (Subtask->WorkingBuffer != NULL) {\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->BlockIo2Token.Event != NULL) {\r
+      gBS->CloseEvent (Subtask->BlockIo2Token.Event);\r
+    }\r
   }\r
+  FreePool (Subtask);\r
 \r
-  return Status;\r
+  return Link;\r
 }\r
 \r
+/**\r
+  The callback for the BlockIo2 ReadBlocksEx/WriteBlocksEx.\r
+  @param  Event                 Event whose notification function is being invoked.\r
+  @param  Context               The pointer to the notification function's context,\r
+                                which points to the DISK_IO_SUBTASK instance.\r
+**/\r
+VOID\r
+EFIAPI\r
+DiskIo2OnReadWriteComplete (\r
+  IN EFI_EVENT            Event,\r
+  IN VOID                 *Context\r
+  )\r
+{\r
+  DISK_IO_SUBTASK       *Subtask;\r
+  DISK_IO2_TASK         *Task;\r
+  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
+  Instance          = Task->Instance;\r
 \r
+  ASSERT (Subtask->Signature  == DISK_IO_SUBTASK_SIGNATURE);\r
+  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
+  }\r
+  RemoveEntryList (&Subtask->Link);\r
+  FreePool (Subtask);\r
+\r
+  if (EFI_ERROR (TransactionStatus) || IsListEmpty (&Task->Subtasks)) {\r
+    if (Task->Token != NULL) {\r
+      //\r
+      // Signal error status once the subtask is failed.\r
+      // Or signal the last status once the last subtask is finished.\r
+      //\r
+      Task->Token->TransactionStatus = TransactionStatus;\r
+      gBS->SignalEvent (Task->Token->Event);\r
+\r
+      //\r
+      // Mark token to NULL\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
-  Read BufferSize bytes from Offset into Buffer.\r
-  Reads may support reads that are not aligned on\r
-  sector boundaries. There are three cases:\r
-    UnderRun - The first byte is not on a sector boundary or the read request is\r
-               less than a sector in length.\r
-    Aligned  - A read of N contiguous sectors.\r
-    OverRun  - The last byte is not on a sector boundary.\r
+  Create the subtask.\r
 \r
-  @param  This                  Protocol instance pointer.\r
-  @param  MediaId               Id of the media, changes every time the media is replaced.\r
-  @param  Offset                The starting byte offset to read from\r
-  @param  BufferSize            Size of Buffer\r
-  @param  Buffer                Buffer containing read data\r
+  @param Write         TRUE: Write request; FALSE: Read request.\r
+  @param Lba           The starting logical block address to read from on the device.\r
+  @param Offset        The starting byte offset to read from the LBA.\r
+  @param Length        The number of bytes to read from the device.\r
+  @param WorkingBuffer The aligned buffer to hold the data for reading or writing.\r
+  @param Buffer        The buffer to hold the data for reading or writing.\r
+  @param Blocking      TRUE: Blocking request; FALSE: Non-blocking request.\r
 \r
-  @retval EFI_SUCCESS           The data was read correctly from the device.\r
-  @retval EFI_DEVICE_ERROR      The device reported an error while performing the read.\r
-  @retval EFI_NO_MEDIA          There is no media in the device.\r
-  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.\r
-  @retval EFI_INVALID_PARAMETER The read request contains device addresses that are not\r
-                                valid for the device.\r
+  @return A pointer to the created subtask.\r
+**/\r
+DISK_IO_SUBTASK *\r
+DiskIoCreateSubtask (\r
+  IN BOOLEAN          Write,\r
+  IN UINT64           Lba,\r
+  IN UINT32           Offset,\r
+  IN UINTN            Length,\r
+  IN VOID             *WorkingBuffer,  OPTIONAL\r
+  IN VOID             *Buffer,\r
+  IN BOOLEAN          Blocking\r
+  )\r
+{\r
+  DISK_IO_SUBTASK       *Subtask;\r
+  EFI_STATUS            Status;\r
+\r
+  Subtask = AllocateZeroPool (sizeof (DISK_IO_SUBTASK));\r
+  if (Subtask == NULL) {\r
+    return NULL;\r
+  }\r
+  Subtask->Signature     = DISK_IO_SUBTASK_SIGNATURE;\r
+  Subtask->Write         = Write;\r
+  Subtask->Lba           = Lba;\r
+  Subtask->Offset        = Offset;\r
+  Subtask->Length        = Length;\r
+  Subtask->WorkingBuffer = WorkingBuffer;\r
+  Subtask->Buffer        = Buffer;\r
+  Subtask->Blocking      = Blocking;\r
+  if (!Blocking) {\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_NOTIFY,\r
+                    DiskIo2OnReadWriteComplete,\r
+                    Subtask,\r
+                    &Subtask->BlockIo2Token.Event\r
+                    );\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (Subtask);\r
+      return NULL;\r
+    }\r
+  }\r
+  DEBUG ((\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
+\r
+  return Subtask;\r
+}\r
 \r
+/**\r
+  Create the subtask list.\r
+\r
+  @param Instance            Pointer to the DISK_IO_PRIVATE_DATA.\r
+  @param Write               TRUE: Write request; FALSE: Read request.\r
+  @param Offset              The starting byte offset to read from the device.\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 buffer for the data.\r
+  @param Blocking            TRUE: Blocking request; FALSE: Non-blocking request.\r
+  @param SharedWorkingBuffer The aligned buffer to hold the data for reading or writing.\r
+  @param Subtasks            The subtask list header.\r
+\r
+  @retval TRUE  The subtask list is created successfully.\r
+  @retval FALSE The subtask list is not created.\r
 **/\r
-EFI_STATUS\r
-EFIAPI\r
-DiskIoReadDisk (\r
-  IN EFI_DISK_IO_PROTOCOL  *This,\r
-  IN UINT32                MediaId,\r
+BOOLEAN\r
+DiskIoCreateSubtaskList (\r
+  IN DISK_IO_PRIVATE_DATA  *Instance,\r
+  IN BOOLEAN               Write,\r
   IN UINT64                Offset,\r
   IN UINTN                 BufferSize,\r
-  OUT VOID                 *Buffer\r
+  IN VOID                  *Buffer,\r
+  IN BOOLEAN               Blocking,\r
+  IN VOID                  *SharedWorkingBuffer,\r
+  IN OUT LIST_ENTRY        *Subtasks\r
   )\r
 {\r
-  EFI_STATUS            Status;\r
-  DISK_IO_PRIVATE_DATA  *Private;\r
-  EFI_BLOCK_IO_PROTOCOL *BlockIo;\r
-  EFI_BLOCK_IO_MEDIA    *Media;\r
   UINT32                BlockSize;\r
+  UINT32                IoAlign;\r
   UINT64                Lba;\r
   UINT64                OverRunLba;\r
   UINT32                UnderRun;\r
   UINT32                OverRun;\r
-  BOOLEAN               TransactionComplete;\r
-  UINTN                 WorkingBufferSize;\r
-  UINT8                 *WorkingBuffer;\r
+  UINT8                 *BufferPtr;\r
   UINTN                 Length;\r
-  UINT8                 *Data;\r
-  UINT8                 *PreData;\r
-  UINTN                 IsBufferAligned;\r
   UINTN                 DataBufferSize;\r
-  BOOLEAN               LastRead;\r
+  DISK_IO_SUBTASK       *Subtask;\r
+  VOID                  *WorkingBuffer;\r
+  LIST_ENTRY            *Link;\r
 \r
-  Private   = DISK_IO_PRIVATE_DATA_FROM_THIS (This);\r
+  DEBUG ((EFI_D_BLKIO, "DiskIo: Create subtasks for task: Offset/BufferSize/Buffer = %016lx/%08x/%08x\n", Offset, BufferSize, Buffer));\r
 \r
-  BlockIo   = Private->BlockIo;\r
-  Media     = BlockIo->Media;\r
-  BlockSize = Media->BlockSize;\r
-\r
-  if (Media->MediaId != MediaId) {\r
-    return EFI_MEDIA_CHANGED;\r
+  BlockSize = Instance->BlockIo->Media->BlockSize;\r
+  IoAlign   = Instance->BlockIo->Media->IoAlign;\r
+  if (IoAlign == 0) {\r
+    IoAlign = 1;\r
   }\r
-\r
-  WorkingBuffer     = Buffer;\r
-  WorkingBufferSize = BufferSize;\r
+  \r
+  Lba       = DivU64x32Remainder (Offset, BlockSize, &UnderRun);\r
+  BufferPtr = (UINT8 *) Buffer;\r
 \r
   //\r
-  // Allocate a temporary buffer for operation\r
+  // Special handling for zero BufferSize\r
   //\r
-  DataBufferSize = BlockSize * DATA_BUFFER_BLOCK_NUM;\r
-\r
-  if (Media->IoAlign > 1) {\r
-    PreData = AllocatePool (DataBufferSize + Media->IoAlign);\r
-    Data    = PreData - ((UINTN) PreData & (Media->IoAlign - 1)) + Media->IoAlign;\r
-  } else {\r
-    PreData = AllocatePool (DataBufferSize);\r
-    Data    = PreData;\r
-  }\r
-\r
-  if (PreData == NULL) {\r
-    return EFI_OUT_OF_RESOURCES;\r
+  if (BufferSize == 0) {\r
+    Subtask = DiskIoCreateSubtask (Write, Lba, UnderRun, 0, NULL, BufferPtr, Blocking);\r
+    if (Subtask == NULL) {\r
+      goto Done;\r
+    }\r
+    InsertTailList (Subtasks, &Subtask->Link);\r
+    return TRUE;\r
   }\r
 \r
-  Lba                 = DivU64x32Remainder (Offset, BlockSize, &UnderRun);\r
-\r
-  Length              = BlockSize - UnderRun;\r
-  TransactionComplete = FALSE;\r
-\r
-  Status              = EFI_SUCCESS;\r
   if (UnderRun != 0) {\r
-    //\r
-    // Offset starts in the middle of an Lba, so read the entire block.\r
-    //\r
-    Status = BlockIo->ReadBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        Lba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
-\r
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\r
+    Length = MIN (BlockSize - UnderRun, BufferSize);\r
+    if (Blocking) {\r
+      WorkingBuffer = SharedWorkingBuffer;\r
+    } else {\r
+      WorkingBuffer = AllocateAlignedPages (EFI_SIZE_TO_PAGES (BlockSize), IoAlign);\r
+      if (WorkingBuffer == NULL) {\r
+        goto Done;\r
+      }\r
+    }\r
+    if (Write) {\r
+      //\r
+      // A half write operation can be splitted to a blocking block-read and half write operation\r
+      // This can simplify the sub task processing logic\r
+      //\r
+      Subtask = DiskIoCreateSubtask (FALSE, Lba, 0, BlockSize, NULL, WorkingBuffer, TRUE);\r
+      if (Subtask == NULL) {\r
+        goto Done;\r
+      }\r
+      InsertTailList (Subtasks, &Subtask->Link);\r
     }\r
 \r
-    if (Length > BufferSize) {\r
-      Length              = BufferSize;\r
-      TransactionComplete = TRUE;\r
+    Subtask = DiskIoCreateSubtask (Write, Lba, UnderRun, Length, WorkingBuffer, BufferPtr, Blocking);\r
+    if (Subtask == NULL) {\r
+      goto Done;\r
     }\r
+    InsertTailList (Subtasks, &Subtask->Link);\r
+  \r
+    BufferPtr  += Length;\r
+    Offset     += Length;\r
+    BufferSize -= Length;\r
+    Lba ++;\r
+  }\r
 \r
-    CopyMem (WorkingBuffer, Data + UnderRun, Length);\r
+  OverRunLba  = Lba + DivU64x32Remainder (BufferSize, BlockSize, &OverRun);\r
+  BufferSize -= OverRun;\r
 \r
-    WorkingBuffer += Length;\r
+  if (OverRun != 0) {\r
+    if (Blocking) {\r
+      WorkingBuffer = SharedWorkingBuffer;\r
+    } else {\r
+      WorkingBuffer = AllocateAlignedPages (EFI_SIZE_TO_PAGES (BlockSize), IoAlign);\r
+      if (WorkingBuffer == NULL) {\r
+        goto Done;\r
+      }\r
+    }\r
+    if (Write) {\r
+      //\r
+      // A half write operation can be splitted to a blocking block-read and half write operation\r
+      // This can simplify the sub task processing logic\r
+      //\r
+      Subtask = DiskIoCreateSubtask (FALSE, OverRunLba, 0, BlockSize, NULL, WorkingBuffer, TRUE);\r
+      if (Subtask == NULL) {\r
+        goto Done;\r
+      }\r
+      InsertTailList (Subtasks, &Subtask->Link);\r
+    }\r
 \r
-    WorkingBufferSize -= Length;\r
-    if (WorkingBufferSize == 0) {\r
+    Subtask = DiskIoCreateSubtask (Write, OverRunLba, 0, OverRun, WorkingBuffer, BufferPtr, Blocking);\r
+    if (Subtask == NULL) {\r
       goto Done;\r
     }\r
-\r
-    Lba += 1;\r
+    InsertTailList (Subtasks, &Subtask->Link);\r
   }\r
-\r
-  OverRunLba = Lba + DivU64x32Remainder (WorkingBufferSize, BlockSize, &OverRun);\r
-\r
-  if (!TransactionComplete && WorkingBufferSize >= BlockSize) {\r
+  \r
+  if (OverRunLba > Lba) {\r
     //\r
     // If the DiskIo maps directly to a BlockIo device do the read.\r
     //\r
-    if (OverRun != 0) {\r
-      WorkingBufferSize -= OverRun;\r
-    }\r
-    //\r
-    // Check buffer alignment\r
-    //\r
-    IsBufferAligned = (UINTN) WorkingBuffer & (UINTN) (Media->IoAlign - 1);\r
-\r
-    if (Media->IoAlign <= 1 || IsBufferAligned == 0) {\r
-      //\r
-      // Alignment is satisfied, so read them together\r
-      //\r
-      Status = BlockIo->ReadBlocks (\r
-                          BlockIo,\r
-                          MediaId,\r
-                          Lba,\r
-                          WorkingBufferSize,\r
-                          WorkingBuffer\r
-                          );\r
-\r
-      if (EFI_ERROR (Status)) {\r
+    if (ALIGN_POINTER (BufferPtr, IoAlign) == BufferPtr) {\r
+      Subtask = DiskIoCreateSubtask (Write, Lba, 0, BufferSize, NULL, BufferPtr, Blocking);\r
+      if (Subtask == NULL) {\r
         goto Done;\r
       }\r
+      InsertTailList (Subtasks, &Subtask->Link);\r
 \r
-      WorkingBuffer += WorkingBufferSize;\r
+      BufferPtr  += BufferSize;\r
+      Offset     += BufferSize;\r
+      BufferSize -= BufferSize;\r
 \r
     } else {\r
-      //\r
-      // Use the allocated buffer instead of the original buffer\r
-      // to avoid alignment issue.\r
-      // Here, the allocated buffer (8-byte align) can satisfy the alignment\r
-      //\r
-      LastRead = FALSE;\r
-      do {\r
-        if (WorkingBufferSize <= DataBufferSize) {\r
+      if (Blocking) {\r
+        //\r
+        // 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
+\r
+          Subtask = DiskIoCreateSubtask (Write, Lba, 0, DataBufferSize, SharedWorkingBuffer, BufferPtr, Blocking);\r
+          if (Subtask == NULL) {\r
+            goto Done;\r
+          }\r
+          InsertTailList (Subtasks, &Subtask->Link);\r
+\r
+          BufferPtr  += DataBufferSize;\r
+          Offset     += DataBufferSize;\r
+          BufferSize -= DataBufferSize;\r
+        }\r
+      } else {\r
+        WorkingBuffer = AllocateAlignedPages (EFI_SIZE_TO_PAGES (BufferSize), IoAlign);\r
+        if (WorkingBuffer == NULL) {\r
           //\r
-          // It is the last calling to readblocks in this loop\r
+          // If there is not enough memory, downgrade to blocking access\r
           //\r
-          DataBufferSize  = WorkingBufferSize;\r
-          LastRead        = TRUE;\r
+          DEBUG ((EFI_D_VERBOSE, "DiskIo: No enough memory so downgrade to blocking access\n"));\r
+          if (!DiskIoCreateSubtaskList (Instance, Write, Offset, BufferSize, BufferPtr, TRUE, SharedWorkingBuffer, Subtasks)) {\r
+            goto Done;\r
+          }\r
+        } else {\r
+          Subtask = DiskIoCreateSubtask (Write, Lba, 0, BufferSize, WorkingBuffer, BufferPtr, Blocking);\r
+          if (Subtask == NULL) {\r
+            goto Done;\r
+          }\r
+          InsertTailList (Subtasks, &Subtask->Link);\r
         }\r
 \r
-        Status = BlockIo->ReadBlocks (\r
-                            BlockIo,\r
-                            MediaId,\r
-                            Lba,\r
-                            DataBufferSize,\r
-                            Data\r
-                            );\r
-        if (EFI_ERROR (Status)) {\r
-          goto Done;\r
-        }\r
-\r
-        CopyMem (WorkingBuffer, Data, DataBufferSize);\r
-        WorkingBufferSize -= DataBufferSize;\r
-        WorkingBuffer += DataBufferSize;\r
-        Lba += DATA_BUFFER_BLOCK_NUM;\r
-      } while (!LastRead);\r
+        BufferPtr  += BufferSize;\r
+        Offset     += BufferSize;\r
+        BufferSize -= BufferSize;\r
+      }\r
     }\r
   }\r
 \r
-  if (!TransactionComplete && OverRun != 0) {\r
-    //\r
-    // Last read is not a complete block.\r
-    //\r
-    Status = BlockIo->ReadBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        OverRunLba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
+  ASSERT (BufferSize == 0);\r
 \r
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\r
-    }\r
-\r
-    CopyMem (WorkingBuffer, Data, OverRun);\r
-  }\r
+  return TRUE;\r
 \r
 Done:\r
-  if (PreData != NULL) {\r
-    FreePool (PreData);\r
+  //\r
+  // Remove all the subtasks.\r
+  //\r
+  for (Link = GetFirstNode (Subtasks); !IsNull (Subtasks, Link); ) {\r
+    Subtask = CR (Link, DISK_IO_SUBTASK, Link, DISK_IO_SUBTASK_SIGNATURE);\r
+    Link = DiskIoDestroySubtask (Instance, Subtask);\r
   }\r
-\r
-  return Status;\r
+  return FALSE;\r
 }\r
 \r
-\r
 /**\r
-  Read BufferSize bytes from Offset into Buffer.\r
-  Writes may require a read modify write to support writes that are not\r
-  aligned on sector boundaries. There are three cases:\r
-    UnderRun - The first byte is not on a sector boundary or the write request\r
-               is less than a sector in length. Read modify write is required.\r
-    Aligned  - A write of N contiguous sectors.\r
-    OverRun  - The last byte is not on a sector boundary. Read modified write\r
-               required.\r
-\r
-  @param  This       Protocol instance pointer.\r
-  @param  MediaId    Id of the media, changes every time the media is replaced.\r
-  @param  Offset     The starting byte offset to read from\r
-  @param  BufferSize Size of Buffer\r
-  @param  Buffer     Buffer containing read data\r
+  Terminate outstanding asynchronous requests to a device.\r
 \r
-  @retval EFI_SUCCESS           The data was written correctly to the device.\r
-  @retval EFI_WRITE_PROTECTED   The device can not be written to.\r
-  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write.\r
-  @retval EFI_NO_MEDIA          There is no media in the device.\r
-  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.\r
-  @retval EFI_INVALID_PARAMETER The write request contains device addresses that are not\r
-                                 valid for the device.\r
+  @param This                   Indicates a pointer to the calling context.\r
 \r
+  @retval EFI_SUCCESS           All outstanding requests were successfully terminated.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the cancel\r
+                                operation.\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
-DiskIoWriteDisk (\r
-  IN EFI_DISK_IO_PROTOCOL  *This,\r
-  IN UINT32                MediaId,\r
-  IN UINT64                Offset,\r
-  IN UINTN                 BufferSize,\r
-  IN VOID                  *Buffer\r
+DiskIo2Cancel (\r
+  IN EFI_DISK_IO2_PROTOCOL *This\r
   )\r
 {\r
-  EFI_STATUS            Status;\r
-  DISK_IO_PRIVATE_DATA  *Private;\r
-  EFI_BLOCK_IO_PROTOCOL *BlockIo;\r
-  EFI_BLOCK_IO_MEDIA    *Media;\r
-  UINT32                BlockSize;\r
-  UINT64                Lba;\r
-  UINT64                OverRunLba;\r
-  UINT32                UnderRun;\r
-  UINT32                OverRun;\r
-  BOOLEAN               TransactionComplete;\r
-  UINTN                 WorkingBufferSize;\r
-  UINT8                 *WorkingBuffer;\r
-  UINTN                 Length;\r
-  UINT8                 *Data;\r
-  UINT8                 *PreData;\r
-  UINTN                 IsBufferAligned;\r
-  UINTN                 DataBufferSize;\r
-  BOOLEAN               LastWrite;\r
+  DISK_IO_PRIVATE_DATA  *Instance;\r
+  DISK_IO2_TASK         *Task;\r
+  LIST_ENTRY            *Link;\r
+  \r
+  Instance = DISK_IO_PRIVATE_DATA_FROM_DISK_IO2 (This);\r
 \r
-  Private   = DISK_IO_PRIVATE_DATA_FROM_THIS (This);\r
+  EfiAcquireLock (&Instance->TaskQueueLock);\r
 \r
-  BlockIo   = Private->BlockIo;\r
-  Media     = BlockIo->Media;\r
-  BlockSize = Media->BlockSize;\r
+  for (Link = GetFirstNode (&Instance->TaskQueue)\r
+    ; !IsNull (&Instance->TaskQueue, Link)\r
+    ; Link = GetNextNode (&Instance->TaskQueue, Link)\r
+    ) {\r
+    Task = CR (Link, DISK_IO2_TASK, Link, DISK_IO2_TASK_SIGNATURE);\r
 \r
-  if (Media->ReadOnly) {\r
-    return EFI_WRITE_PROTECTED;\r
+    if (Task->Token != NULL) {\r
+      Task->Token->TransactionStatus = EFI_ABORTED;\r
+      gBS->SignalEvent (Task->Token->Event);\r
+      //\r
+      // Set Token to NULL so that the further BlockIo2 responses will be ignored\r
+      //\r
+      Task->Token = NULL;\r
+    }\r
   }\r
 \r
-  if (Media->MediaId != MediaId) {\r
-    return EFI_MEDIA_CHANGED;\r
-  }\r
+  EfiReleaseLock (&Instance->TaskQueueLock);\r
 \r
-  DataBufferSize = BlockSize * DATA_BUFFER_BLOCK_NUM;\r
+  return EFI_SUCCESS;\r
+}\r
 \r
-  if (Media->IoAlign > 1) {\r
-    PreData = AllocatePool (DataBufferSize + Media->IoAlign);\r
-    Data    = PreData - ((UINTN) PreData & (Media->IoAlign - 1)) + Media->IoAlign;\r
-  } else {\r
-    PreData = AllocatePool (DataBufferSize);\r
-    Data    = PreData;\r
-  }\r
+/**\r
+  Common routine to access the disk.\r
+\r
+  @param Instance    Pointer to the DISK_IO_PRIVATE_DATA.\r
+  @param Write       TRUE: Write operation; FALSE: Read operation.\r
+  @param MediaId     ID of the medium to access.\r
+  @param Offset      The starting byte offset on the logical block I/O device to access.\r
+  @param Token       A pointer to the token associated with the transaction.\r
+                     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
+**/\r
+EFI_STATUS\r
+DiskIo2ReadWriteDisk (\r
+  IN DISK_IO_PRIVATE_DATA     *Instance,\r
+  IN BOOLEAN                  Write,\r
+  IN UINT32                   MediaId,\r
+  IN UINT64                   Offset,\r
+  IN EFI_DISK_IO2_TOKEN       *Token,\r
+  IN UINTN                    BufferSize,\r
+  IN UINT8                    *Buffer\r
+  )\r
+{\r
+  EFI_STATUS             Status;\r
+  EFI_BLOCK_IO_PROTOCOL  *BlockIo;\r
+  EFI_BLOCK_IO2_PROTOCOL *BlockIo2;\r
+  EFI_BLOCK_IO_MEDIA     *Media;\r
+  LIST_ENTRY             *Link;\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
+  LIST_ENTRY             *SubtasksPtr;\r
+\r
+  Task      = NULL;\r
+  BlockIo   = Instance->BlockIo;\r
+  BlockIo2  = Instance->BlockIo2;\r
+  Media     = BlockIo->Media;\r
+  Status    = EFI_SUCCESS;\r
+  Blocking  = ((Token == NULL) || (Token->Event == NULL));\r
 \r
-  if (PreData == NULL) {\r
-    return EFI_OUT_OF_RESOURCES;\r
+  if (Media->MediaId != MediaId) {\r
+    return EFI_MEDIA_CHANGED;\r
   }\r
-\r
-  WorkingBuffer       = Buffer;\r
-  WorkingBufferSize   = BufferSize;\r
-\r
-  Lba                 = DivU64x32Remainder (Offset, BlockSize, &UnderRun);\r
-\r
-  Length              = BlockSize - UnderRun;\r
-  TransactionComplete = FALSE;\r
-\r
-  Status              = EFI_SUCCESS;\r
-  if (UnderRun != 0) {\r
+  \r
+  if (Write && Media->ReadOnly) {\r
+    return EFI_WRITE_PROTECTED;\r
+  }\r
+  \r
+  if (Blocking) {\r
     //\r
-    // Offset starts in the middle of an Lba, so do read modify write.\r
+    // Wait till pending async task is completed.\r
     //\r
-    Status = BlockIo->ReadBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        Lba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
+    do {\r
+      EfiAcquireLock (&Instance->TaskQueueLock);\r
+      TaskQueueEmpty = IsListEmpty (&Instance->TaskQueue);\r
+      EfiReleaseLock (&Instance->TaskQueueLock);\r
+    } while (!TaskQueueEmpty);\r
 \r
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\r
-    }\r
-\r
-    if (Length > BufferSize) {\r
-      Length              = BufferSize;\r
-      TransactionComplete = TRUE;\r
+    SubtasksPtr = &Subtasks;\r
+  } else {\r
+    Task = AllocatePool (sizeof (DISK_IO2_TASK));\r
+    if (Task == NULL) {\r
+      return EFI_OUT_OF_RESOURCES;\r
     }\r
+    \r
+    EfiAcquireLock (&Instance->TaskQueueLock);\r
+    InsertTailList (&Instance->TaskQueue, &Task->Link);\r
+    EfiReleaseLock (&Instance->TaskQueueLock);\r
 \r
-    CopyMem (Data + UnderRun, WorkingBuffer, Length);\r
+    Task->Signature = DISK_IO2_TASK_SIGNATURE;\r
+    Task->Instance  = Instance;\r
+    Task->Token     = Token;\r
 \r
-    Status = BlockIo->WriteBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        Lba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\r
-    }\r
+    SubtasksPtr = &Task->Subtasks;\r
+  }\r
 \r
-    WorkingBuffer += Length;\r
-    WorkingBufferSize -= Length;\r
-    if (WorkingBufferSize == 0) {\r
-      goto Done;\r
+  InitializeListHead (SubtasksPtr);\r
+  if (!DiskIoCreateSubtaskList (Instance, Write, Offset, BufferSize, Buffer, Blocking, Instance->SharedWorkingBuffer, SubtasksPtr)) {\r
+    if (Task != NULL) {\r
+      FreePool (Task);\r
     }\r
-\r
-    Lba += 1;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
+  ASSERT (!IsListEmpty (SubtasksPtr));\r
 \r
-  OverRunLba = Lba + DivU64x32Remainder (WorkingBufferSize, BlockSize, &OverRun);\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
 \r
-  if (!TransactionComplete && WorkingBufferSize >= BlockSize) {\r
-    //\r
-    // If the DiskIo maps directly to a BlockIo device do the write.\r
-    //\r
-    if (OverRun != 0) {\r
-      WorkingBufferSize -= OverRun;\r
+    if (!Subtask->Blocking) {\r
+      Subtask->Task = Task;\r
     }\r
-    //\r
-    // Check buffer alignment\r
-    //\r
-    IsBufferAligned = (UINTN) WorkingBuffer & (UINTN) (Media->IoAlign - 1);\r
 \r
-    if (Media->IoAlign <= 1 || IsBufferAligned == 0) {\r
+    if (Subtask->Write) {\r
       //\r
-      // Alignment is satisfied, so write them together\r
+      // Write\r
       //\r
-      Status = BlockIo->WriteBlocks (\r
-                          BlockIo,\r
-                          MediaId,\r
-                          Lba,\r
-                          WorkingBufferSize,\r
-                          WorkingBuffer\r
-                          );\r
-\r
-      if (EFI_ERROR (Status)) {\r
-        goto Done;\r
+      if (Subtask->WorkingBuffer != NULL) {\r
+        //\r
+        // A sub task before this one should be a block read operation, causing the WorkingBuffer filled with the entire one block data.\r
+        //\r
+        CopyMem (Subtask->WorkingBuffer + Subtask->Offset, Subtask->Buffer, Subtask->Length);\r
       }\r
 \r
-      WorkingBuffer += WorkingBufferSize;\r
+      if (Subtask->Blocking) {\r
+        Status = BlockIo->WriteBlocks (\r
+                            BlockIo,\r
+                            MediaId,\r
+                            Subtask->Lba,\r
+                            Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                            (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
+                            );\r
+      } else {\r
+        Status = BlockIo2->WriteBlocksEx (\r
+                             BlockIo2,\r
+                             MediaId,\r
+                             Subtask->Lba,\r
+                             &Subtask->BlockIo2Token,\r
+                             Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                             (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
+                             );\r
+      }\r
 \r
     } else {\r
       //\r
-      // The buffer parameter is not aligned with the request\r
-      // So use the allocated instead.\r
-      // It can fit almost all the cases.\r
+      // Read\r
       //\r
-      LastWrite = FALSE;\r
-      do {\r
-        if (WorkingBufferSize <= DataBufferSize) {\r
-          //\r
-          // It is the last calling to writeblocks in this loop\r
-          //\r
-          DataBufferSize  = WorkingBufferSize;\r
-          LastWrite       = TRUE;\r
-        }\r
-\r
-        CopyMem (Data, WorkingBuffer, DataBufferSize);\r
-        Status = BlockIo->WriteBlocks (\r
+      if (Subtask->Blocking) {\r
+        Status = BlockIo->ReadBlocks (\r
                             BlockIo,\r
                             MediaId,\r
-                            Lba,\r
-                            DataBufferSize,\r
-                            Data\r
+                            Subtask->Lba,\r
+                            Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                            (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
                             );\r
-        if (EFI_ERROR (Status)) {\r
-          goto Done;\r
+        if (!EFI_ERROR (Status) && (Subtask->WorkingBuffer != NULL)) {\r
+          CopyMem (Subtask->Buffer, Subtask->WorkingBuffer + Subtask->Offset, Subtask->Length);\r
         }\r
-\r
-        WorkingBufferSize -= DataBufferSize;\r
-        WorkingBuffer += DataBufferSize;\r
-        Lba += DATA_BUFFER_BLOCK_NUM;\r
-      } while (!LastWrite);\r
+      } else {\r
+        Status = BlockIo2->ReadBlocksEx (\r
+                             BlockIo2,\r
+                             MediaId,\r
+                             Subtask->Lba,\r
+                             &Subtask->BlockIo2Token,\r
+                             Subtask->Length < Media->BlockSize ? Media->BlockSize : Subtask->Length,\r
+                             (Subtask->WorkingBuffer != NULL) ? Subtask->WorkingBuffer : Subtask->Buffer\r
+                             );\r
+      }\r
+    }\r
+    \r
+    if (EFI_ERROR (Status)) {\r
+      break;\r
     }\r
   }\r
+  \r
+  gBS->RaiseTPL (TPL_NOTIFY);\r
 \r
-  if (!TransactionComplete && OverRun != 0) {\r
-    //\r
-    // Last bit is not a complete block, so do a read modify write.\r
-    //\r
-    Status = BlockIo->ReadBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        OverRunLba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
+  //\r
+  // Remove all the remaining subtasks when failure.\r
+  // 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
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\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
     }\r
+  }\r
 \r
-    CopyMem (Data, WorkingBuffer, OverRun);\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
+  //\r
+  if (!Blocking && IsListEmpty (SubtasksPtr)) {\r
+    EfiAcquireLock (&Instance->TaskQueueLock);\r
+    RemoveEntryList (&Task->Link);\r
+    EfiReleaseLock (&Instance->TaskQueueLock);\r
 \r
-    Status = BlockIo->WriteBlocks (\r
-                        BlockIo,\r
-                        MediaId,\r
-                        OverRunLba,\r
-                        BlockSize,\r
-                        Data\r
-                        );\r
-    if (EFI_ERROR (Status)) {\r
-      goto Done;\r
+    if (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
+      //\r
+      DEBUG ((EFI_D_VERBOSE, "DiskIo: Non-blocking request was downgraded to blocking request, signal event directly.\n"));\r
+      Task->Token->TransactionStatus = Status;\r
+      gBS->SignalEvent (Task->Token->Event);\r
     }\r
+\r
+    FreePool (Task);\r
   }\r
 \r
-Done:\r
-  if (PreData != NULL) {\r
-    FreePool (PreData);\r
+  gBS->RestoreTPL (OldTpl);\r
+\r
+  return Status;\r
+}\r
+\r
+/**\r
+  Reads a specified number of bytes from a device.\r
+\r
+  @param This                   Indicates a pointer to the calling context.\r
+  @param MediaId                ID of the medium to be read.\r
+  @param Offset                 The starting byte offset on the logical block I/O device to read from.\r
+  @param Token                  A pointer to the token associated with the transaction.\r
+                                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
+\r
+  @retval EFI_SUCCESS           If Event is NULL (blocking I/O): The data was read correctly from the device.\r
+                                If Event is not NULL (asynchronous I/O): The request was successfully queued for processing.\r
+                                                                         Event will be signaled upon completion.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write.\r
+  @retval EFI_NO_MEDIA          There is no medium in the device.\r
+  @retval EFI_MEDIA_CHNAGED     The MediaId is not for the current medium.\r
+  @retval EFI_INVALID_PARAMETER The read request contains device addresses that are not valid for the device.\r
+  @retval EFI_OUT_OF_RESOURCES  The request could not be completed due to a lack of resources.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DiskIo2ReadDiskEx (\r
+  IN EFI_DISK_IO2_PROTOCOL        *This,\r
+  IN UINT32                       MediaId,\r
+  IN UINT64                       Offset,\r
+  IN OUT EFI_DISK_IO2_TOKEN       *Token,\r
+  IN UINTN                        BufferSize,\r
+  OUT VOID                        *Buffer\r
+  )\r
+{\r
+  return DiskIo2ReadWriteDisk (\r
+           DISK_IO_PRIVATE_DATA_FROM_DISK_IO2 (This),\r
+           FALSE, MediaId, Offset, Token, BufferSize, (UINT8 *) Buffer\r
+           );\r
+}\r
+\r
+/**\r
+  Writes a specified number of bytes to a device.\r
+\r
+  @param This        Indicates a pointer to the calling context.\r
+  @param MediaId     ID of the medium to be written.\r
+  @param Offset      The starting byte offset on the logical block I/O device to write to.\r
+  @param Token       A pointer to the token associated with the transaction.\r
+                     If this field is NULL, synchronous/blocking IO is performed.\r
+  @param BufferSize  The size in bytes of Buffer. The number of bytes to write to the device.\r
+  @param Buffer      A pointer to the buffer containing the data to be written.\r
+\r
+  @retval EFI_SUCCESS           If Event is NULL (blocking I/O): The data was written correctly to the device.\r
+                                If Event is not NULL (asynchronous I/O): The request was successfully queued for processing.\r
+                                                                         Event will be signaled upon completion.\r
+  @retval EFI_WRITE_PROTECTED   The device cannot be written to.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write operation.\r
+  @retval EFI_NO_MEDIA          There is no medium in the device.\r
+  @retval EFI_MEDIA_CHNAGED     The MediaId is not for the current medium.\r
+  @retval EFI_INVALID_PARAMETER The write request contains device addresses that are not valid for the device.\r
+  @retval EFI_OUT_OF_RESOURCES  The request could not be completed due to a lack of resources.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DiskIo2WriteDiskEx (\r
+  IN EFI_DISK_IO2_PROTOCOL        *This,\r
+  IN UINT32                       MediaId,\r
+  IN UINT64                       Offset,\r
+  IN OUT EFI_DISK_IO2_TOKEN       *Token,\r
+  IN UINTN                        BufferSize,\r
+  IN VOID                         *Buffer\r
+  )\r
+{\r
+  return DiskIo2ReadWriteDisk (\r
+           DISK_IO_PRIVATE_DATA_FROM_DISK_IO2 (This),\r
+           TRUE, MediaId, Offset, Token, BufferSize, (UINT8 *) Buffer\r
+           );\r
+}\r
+\r
+/**\r
+  The callback for the BlockIo2 FlushBlocksEx.\r
+  @param  Event                 Event whose notification function is being invoked.\r
+  @param  Context               The pointer to the notification function's context,\r
+                                which points to the DISK_IO2_FLUSH_TASK instance.\r
+**/\r
+VOID\r
+EFIAPI\r
+DiskIo2OnFlushComplete (\r
+  IN EFI_EVENT                 Event,\r
+  IN VOID                      *Context\r
+  )\r
+{\r
+  DISK_IO2_FLUSH_TASK             *Task;\r
+\r
+  gBS->CloseEvent (Event);\r
+\r
+  Task = (DISK_IO2_FLUSH_TASK *) Context;\r
+  ASSERT (Task->Signature == DISK_IO2_FLUSH_TASK_SIGNATURE);\r
+  Task->Token->TransactionStatus = Task->BlockIo2Token.TransactionStatus;\r
+  gBS->SignalEvent (Task->Token->Event);\r
+}\r
+\r
+/**\r
+  Flushes all modified data to the physical device.\r
+\r
+  @param This        Indicates a pointer to the calling context.\r
+  @param Token       A pointer to the token associated with the transaction.\r
+                     If this field is NULL, synchronous/blocking IO is performed.\r
+\r
+  @retval EFI_SUCCESS           If Event is NULL (blocking I/O): The data was flushed successfully to the device.\r
+                                If Event is not NULL (asynchronous I/O): The request was successfully queued for processing.\r
+                                                                         Event will be signaled upon completion.\r
+  @retval EFI_WRITE_PROTECTED   The device cannot be written to.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write operation.\r
+  @retval EFI_NO_MEDIA          There is no medium in the device.\r
+  @retval EFI_OUT_OF_RESOURCES  The request could not be completed due to a lack of resources.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DiskIo2FlushDiskEx (\r
+  IN EFI_DISK_IO2_PROTOCOL        *This,\r
+  IN OUT EFI_DISK_IO2_TOKEN       *Token\r
+  )\r
+{\r
+  EFI_STATUS                      Status;\r
+  DISK_IO2_FLUSH_TASK             *Task;\r
+  DISK_IO_PRIVATE_DATA            *Private;\r
+\r
+  Private = DISK_IO_PRIVATE_DATA_FROM_DISK_IO2 (This);\r
+\r
+  if ((Token != NULL) && (Token->Event != NULL)) {\r
+    Task = AllocatePool (sizeof (DISK_IO2_FLUSH_TASK));\r
+    if (Task == NULL) {\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
+\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_CALLBACK,\r
+                    DiskIo2OnFlushComplete,\r
+                    Task,\r
+                    &Task->BlockIo2Token.Event\r
+                    );\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (Task);\r
+      return Status;\r
+    }\r
+    Task->Signature = DISK_IO2_FLUSH_TASK_SIGNATURE;\r
+    Status = Private->BlockIo2->FlushBlocksEx (Private->BlockIo2, &Task->BlockIo2Token);\r
+    if (EFI_ERROR (Status)) {\r
+      gBS->CloseEvent (Task->BlockIo2Token.Event);\r
+      FreePool (Task);\r
+    }\r
+  } else {\r
+    Status = Private->BlockIo2->FlushBlocksEx (Private->BlockIo2, NULL);\r
   }\r
 \r
   return Status;\r
 }\r
 \r
+/**\r
+  Read BufferSize bytes from Offset into Buffer.\r
+  Reads may support reads that are not aligned on\r
+  sector boundaries. There are three cases:\r
+    UnderRun - The first byte is not on a sector boundary or the read request is\r
+               less than a sector in length.\r
+    Aligned  - A read of N contiguous sectors.\r
+    OverRun  - The last byte is not on a sector boundary.\r
+\r
+  @param  This                  Protocol instance pointer.\r
+  @param  MediaId               Id of the media, changes every time the media is replaced.\r
+  @param  Offset                The starting byte offset to read from\r
+  @param  BufferSize            Size of Buffer\r
+  @param  Buffer                Buffer containing read data\r
+\r
+  @retval EFI_SUCCESS           The data was read correctly from the device.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the read.\r
+  @retval EFI_NO_MEDIA          There is no media in the device.\r
+  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.\r
+  @retval EFI_INVALID_PARAMETER The read request contains device addresses that are not\r
+                                valid for the device.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DiskIoReadDisk (\r
+  IN EFI_DISK_IO_PROTOCOL  *This,\r
+  IN UINT32                MediaId,\r
+  IN UINT64                Offset,\r
+  IN UINTN                 BufferSize,\r
+  OUT VOID                 *Buffer\r
+  )\r
+{\r
+  return DiskIo2ReadWriteDisk (\r
+           DISK_IO_PRIVATE_DATA_FROM_DISK_IO (This),\r
+           FALSE, MediaId, Offset, NULL, BufferSize, (UINT8 *) Buffer\r
+           );\r
+}\r
+\r
+\r
+/**\r
+  Writes BufferSize bytes from Buffer into Offset.\r
+  Writes may require a read modify write to support writes that are not\r
+  aligned on sector boundaries. There are three cases:\r
+    UnderRun - The first byte is not on a sector boundary or the write request\r
+               is less than a sector in length. Read modify write is required.\r
+    Aligned  - A write of N contiguous sectors.\r
+    OverRun  - The last byte is not on a sector boundary. Read modified write\r
+               required.\r
+\r
+  @param  This       Protocol instance pointer.\r
+  @param  MediaId    Id of the media, changes every time the media is replaced.\r
+  @param  Offset     The starting byte offset to read from\r
+  @param  BufferSize Size of Buffer\r
+  @param  Buffer     Buffer containing read data\r
+\r
+  @retval EFI_SUCCESS           The data was written correctly to the device.\r
+  @retval EFI_WRITE_PROTECTED   The device can not be written to.\r
+  @retval EFI_DEVICE_ERROR      The device reported an error while performing the write.\r
+  @retval EFI_NO_MEDIA          There is no media in the device.\r
+  @retval EFI_MEDIA_CHNAGED     The MediaId does not matched the current device.\r
+  @retval EFI_INVALID_PARAMETER The write request contains device addresses that are not\r
+                                 valid for the device.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DiskIoWriteDisk (\r
+  IN EFI_DISK_IO_PROTOCOL  *This,\r
+  IN UINT32                MediaId,\r
+  IN UINT64                Offset,\r
+  IN UINTN                 BufferSize,\r
+  IN VOID                  *Buffer\r
+  )\r
+{\r
+  return DiskIo2ReadWriteDisk (\r
+           DISK_IO_PRIVATE_DATA_FROM_DISK_IO (This),\r
+           TRUE, MediaId, Offset, NULL, BufferSize, (UINT8 *) Buffer\r
+           );\r
+}\r
 \r
 /**\r
   The user Entry Point for module DiskIo. The user code starts with this function.\r
@@ -720,18 +1243,15 @@ InitializeDiskIo (
   //\r
   // Install driver model protocol(s).\r
   //\r
-  Status = EfiLibInstallAllDriverProtocols (\r
+  Status = EfiLibInstallDriverBindingComponentName2 (\r
              ImageHandle,\r
              SystemTable,\r
              &gDiskIoDriverBinding,\r
              ImageHandle,\r
              &gDiskIoComponentName,\r
-             NULL,\r
-             NULL\r
+             &gDiskIoComponentName2\r
              );\r
   ASSERT_EFI_ERROR (Status);\r
 \r
-\r
   return Status;\r
 }\r
-\r