From 58abc1e5f9c374f2b26634b7f41def82ae3233f6 Mon Sep 17 00:00:00 2001
From: jiangdongxu <jiangdongxu1@huawei.com>
Date: Sat, 7 Sep 2024 17:30:28 +0800
Subject: [PATCH 3/3] VirtioBlk: split large IO according to segment_size_max

When the VirtioBlk device is initialized, the value of SegmentSizeMax
is obtained based on the feature capability. Then delivere the requests
 based on the value of SegmentSizeMax.

Signed-off-by: jiangdongxu <jiangdongxu1@huawei.com>
---
 MdePkg/Include/Protocol/BlockIo.h |  10 ++
 OvmfPkg/VirtioBlkDxe/VirtioBlk.c  | 147 +++++++++++++++++++++---------
 2 files changed, 116 insertions(+), 41 deletions(-)

diff --git a/MdePkg/Include/Protocol/BlockIo.h b/MdePkg/Include/Protocol/BlockIo.h
index 3bd7688..30b20d0 100644
--- a/MdePkg/Include/Protocol/BlockIo.h
+++ b/MdePkg/Include/Protocol/BlockIo.h
@@ -197,6 +197,16 @@ typedef struct {
   /// granularity as a number of logical blocks.
   ///
   UINT32 OptimalTransferLengthGranularity;
+
+  ///
+  /// Maximum size of any single segment
+  ///
+  UINT32  MaxSegmentSize;
+
+  ///
+  /// Maximum number of segments in a request
+  ///
+  UINT32  MaxSegments;
 } EFI_BLOCK_IO_MEDIA;
 
 #define EFI_BLOCK_IO_PROTOCOL_REVISION  0x00010000
diff --git a/OvmfPkg/VirtioBlkDxe/VirtioBlk.c b/OvmfPkg/VirtioBlkDxe/VirtioBlk.c
index 70fc8dd..d13cb37 100644
--- a/OvmfPkg/VirtioBlkDxe/VirtioBlk.c
+++ b/OvmfPkg/VirtioBlkDxe/VirtioBlk.c
@@ -31,6 +31,8 @@
 #define MAX_RETRY_TIMES 1000
 #define DEVICE_WAIT_INTVL 1000
 
+#define DEFAULT_MAX_SEGMENTS 32
+
 /**
 
   Convenience macros to read and write region 0 IO space elements of the
@@ -457,6 +459,67 @@ FreeHostStatusBuffer:
   return Status;
 }
 
+STATIC
+EFI_STATUS
+EFIAPI
+VirtioBlkReadWriteBlocks (
+  IN EFI_BLOCK_IO_PROTOCOL *This,
+  IN UINT32                MediaId,
+  IN EFI_LBA               Lba,
+  IN UINTN                 BufferSize,
+  IN OUT VOID              *Buffer,
+  IN BOOLEAN               RequestIsWrite
+  )
+{
+  VBLK_DEV   *Dev;
+  EFI_STATUS Status;
+  UINT32 SizeMax;
+
+  if (BufferSize == 0) {
+    return EFI_SUCCESS;
+  }
+
+  Dev = VIRTIO_BLK_FROM_BLOCK_IO (This);
+  Status = VerifyReadWriteRequest (
+             &Dev->BlockIoMedia,
+             Lba,
+             BufferSize,
+             RequestIsWrite
+             );
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  SizeMax = Dev->BlockIoMedia.MaxSegmentSize;
+  while (BufferSize >= SizeMax) {
+    Status = SynchronousRequest (
+               Dev,
+               Lba,
+               SizeMax,
+               Buffer,
+               RequestIsWrite
+               );
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+
+    Lba += SizeMax / Dev->BlockIoMedia.BlockSize;
+    BufferSize -= SizeMax;
+    Buffer = (CHAR8 *)Buffer + SizeMax;
+  }
+
+  if (BufferSize == 0) {
+    return EFI_SUCCESS;
+  }
+
+  return SynchronousRequest (
+           Dev,
+           Lba,
+           BufferSize,
+           Buffer,
+           RequestIsWrite
+           );
+}
 
 /**
 
@@ -486,30 +549,13 @@ VirtioBlkReadBlocks (
   OUT VOID                  *Buffer
   )
 {
-  VBLK_DEV   *Dev;
-  EFI_STATUS Status;
-
-  if (BufferSize == 0) {
-    return EFI_SUCCESS;
-  }
-
-  Dev = VIRTIO_BLK_FROM_BLOCK_IO (This);
-  Status = VerifyReadWriteRequest (
-             &Dev->BlockIoMedia,
-             Lba,
-             BufferSize,
-             FALSE               // RequestIsWrite
-             );
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  return SynchronousRequest (
-           Dev,
+  return VirtioBlkReadWriteBlocks(
+           This,
+           MediaId,
            Lba,
            BufferSize,
            Buffer,
-           FALSE       // RequestIsWrite
+           FALSE     // RequestIsRead
            );
 }
 
@@ -541,26 +587,9 @@ VirtioBlkWriteBlocks (
   IN VOID                  *Buffer
   )
 {
-  VBLK_DEV   *Dev;
-  EFI_STATUS Status;
-
-  if (BufferSize == 0) {
-    return EFI_SUCCESS;
-  }
-
-  Dev = VIRTIO_BLK_FROM_BLOCK_IO (This);
-  Status = VerifyReadWriteRequest (
-             &Dev->BlockIoMedia,
-             Lba,
-             BufferSize,
-             TRUE                // RequestIsWrite
-             );
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  return SynchronousRequest (
-           Dev,
+  return VirtioBlkReadWriteBlocks(
+           This,
+           MediaId,
            Lba,
            BufferSize,
            Buffer,
@@ -718,6 +747,8 @@ VirtioBlkInit (
   UINT8      NextDevStat;
   EFI_STATUS Status;
 
+  UINT32     MaxSegmentSize;
+  UINT32     MaxSegments;
   UINT64     Features;
   UINT64     NumSectors;
   UINT32     BlockSize;
@@ -814,6 +845,36 @@ VirtioBlkInit (
     BlockSize = 512;
   }
 
+  if (Features & VIRTIO_BLK_F_SIZE_MAX) {
+    Status = VIRTIO_CFG_READ (Dev, SizeMax, &MaxSegmentSize);
+    if (EFI_ERROR (Status)) {
+      goto Failed;
+    }
+    if (MaxSegmentSize == 0) {
+      //
+      // We need at least one 4KB page.
+      //
+      MaxSegmentSize = SIZE_4KB;
+    }
+  } else {
+    MaxSegmentSize = SIZE_512KB;
+  }
+ 
+  if (Features & VIRTIO_BLK_F_SEG_MAX) {
+    Status = VIRTIO_CFG_READ (Dev, SegMax, &MaxSegments);
+    if (EFI_ERROR (Status)) {
+      goto Failed;
+    }
+    if (MaxSegments == 0) {
+      //
+      // We need at least one SG element, whatever they say.
+      //
+      MaxSegments = 1;
+    }
+  } else {
+    MaxSegments = DEFAULT_MAX_SEGMENTS;
+  }
+
   if (Features & VIRTIO_BLK_F_TOPOLOGY) {
     Status = VIRTIO_CFG_READ (Dev, Topology.PhysicalBlockExp,
                &PhysicalBlockExp);
@@ -949,6 +1010,8 @@ VirtioBlkInit (
   Dev->BlockIoMedia.ReadOnly         = (BOOLEAN) ((Features & VIRTIO_BLK_F_RO) != 0);
   Dev->BlockIoMedia.WriteCaching     = (BOOLEAN) ((Features & VIRTIO_BLK_F_FLUSH) != 0);
   Dev->BlockIoMedia.BlockSize        = BlockSize;
+  Dev->BlockIoMedia.MaxSegments      = MaxSegments;
+  Dev->BlockIoMedia.MaxSegmentSize   = MaxSegmentSize;
   Dev->BlockIoMedia.IoAlign          = 0;
   Dev->BlockIoMedia.LastBlock        = DivU64x32 (NumSectors,
                                          BlockSize / 512) - 1;
@@ -956,6 +1019,8 @@ VirtioBlkInit (
   DEBUG ((DEBUG_INFO, "%a: LbaSize=0x%x[B] NumBlocks=0x%Lx[Lba]\n",
     __FUNCTION__, Dev->BlockIoMedia.BlockSize,
     Dev->BlockIoMedia.LastBlock + 1));
+  DEBUG ((DEBUG_INFO, "%a: MaxSegments=0x%x[B] MaxSegmentSize=0x%x[B]\n",
+    __FUNCTION__, Dev->BlockIoMedia.MaxSegments, Dev->BlockIoMedia.MaxSegmentSize));
 
   if (Features & VIRTIO_BLK_F_TOPOLOGY) {
     Dev->BlockIo.Revision = EFI_BLOCK_IO_PROTOCOL_REVISION3;
-- 
2.46.0.windows.1

