fix: Add buffered transport API and read optimizations

13584452567 · 13584452567 · commit e6595169eb79 · 2026-03-07T13:06:15.000+08:00
Introduce IFastbootBufferedTransport with ReadInto to allow callers to supply buffers and avoid per-read allocations. Implement the buffered API in TcpTransport, UdpTransport and UsbDevice, refactor their Read methods to use ReadInto, and add validation for offsets/lengths. Increase UDP handshake retry budget (min 5 attempts) to reduce fragile CI timing failures. Optimize InternalUploadData to rent a shared buffer from ArrayPool and use the buffered transport fast path when available, reducing allocations and improving throughput; ensure proper return of pooled buffers and EOF/error handling.
diff --git a/FirmwareKit.Comm.Fastboot/Backend/Network/IFastbootTransport.cs b/FirmwareKit.Comm.Fastboot/Backend/Network/IFastbootTransport.cs
@@ -7,4 +7,13 @@ public interface IFastbootTransport : IDisposable
     long Write(byte[] data, int length);
 }
 
+/// <summary>
+/// Optional transport extension for reading directly into caller-provided buffers.
+/// Implement this to avoid per-read byte[] allocations on hot paths.
+/// </summary>
+public interface IFastbootBufferedTransport : IFastbootTransport
+{
+    int ReadInto(byte[] buffer, int offset, int length);
+}
+
 
diff --git a/FirmwareKit.Comm.Fastboot/Backend/Network/TcpTransport.cs b/FirmwareKit.Comm.Fastboot/Backend/Network/TcpTransport.cs
@@ -4,7 +4,7 @@
 
 namespace FirmwareKit.Comm.Fastboot.Network;
 
-public class TcpTransport : IFastbootTransport
+public class TcpTransport : IFastbootBufferedTransport
 {
     private const int DefaultIoTimeoutMs = 30000;
     private readonly TcpClient _client = new();
@@ -75,6 +75,25 @@ private int ReadFully(byte[] buffer, int offset, int length)
 
     public byte[] Read(int length)
     {
+        if (length <= 0) return Array.Empty<byte>();
+
+        byte[] dataBuffer = new byte[length];
+        int actuallyRead = ReadInto(dataBuffer, 0, length);
+        if (actuallyRead < length)
+        {
+            Array.Resize(ref dataBuffer, actuallyRead);
+        }
+        return dataBuffer;
+    }
+
+    public int ReadInto(byte[] buffer, int offset, int length)
+    {
+        if (length <= 0) return 0;
+        if (offset < 0 || length < 0 || offset + length > buffer.Length)
+        {
+            throw new ArgumentOutOfRangeException(nameof(length));
+        }
+
         if (_messageBytesLeft == 0)
         {
             byte[] lenBuffer = new byte[8];
@@ -86,14 +105,9 @@ public byte[] Read(int length)
         }
 
         int toRead = (int)Math.Min(length, _messageBytesLeft);
-        byte[] dataBuffer = new byte[toRead];
-        int actuallyRead = ReadFully(dataBuffer, 0, toRead);
-        if (actuallyRead < toRead)
-        {
-            Array.Resize(ref dataBuffer, actuallyRead);
-        }
+        int actuallyRead = ReadFully(buffer, offset, toRead);
         _messageBytesLeft -= actuallyRead;
-        return dataBuffer;
+        return actuallyRead;
     }
 
     public long Write(byte[] data, int length)
diff --git a/FirmwareKit.Comm.Fastboot/Backend/Network/UdpTransport.cs b/FirmwareKit.Comm.Fastboot/Backend/Network/UdpTransport.cs
@@ -7,7 +7,7 @@ namespace FirmwareKit.Comm.Fastboot.Network;
 /// Fastboot over UDP Transport
 /// Fully implements the AOSP Fastboot over Network protocol (headers, sequence numbers, handshake)
 /// </summary>
-public class UdpTransport : IFastbootTransport
+public class UdpTransport : IFastbootBufferedTransport
 {
     private readonly UdpClient _client;
     private readonly IPEndPoint _endpoint;
@@ -55,13 +55,18 @@ private void InitializeProtocol()
         _client.Client.ReceiveTimeout = _timeoutMs;
         _client.Client.SendTimeout = _timeoutMs;
 
-        byte[] response = SendSinglePacket(PacketId.DeviceQuery, 0, PacketFlag.None, [], 0, 0, _maxTransmissionAttempts, out _);
+        // Handshake runs at transport creation and is sensitive to scheduling jitter
+        // in constrained CI environments. Use a slightly larger retry budget here
+        // without changing steady-state transfer behavior.
+        int initAttempts = Math.Max(_maxTransmissionAttempts, 5);
+
+        byte[] response = SendSinglePacket(PacketId.DeviceQuery, 0, PacketFlag.None, [], 0, 0, initAttempts, out _);
         if (response.Length < 2) throw new Exception("Invalid query response from target.");
         _sequence = BinaryPrimitives.ReadUInt16BigEndian(response.AsSpan(0, 2));
         byte[] initData = new byte[4];
         BinaryPrimitives.WriteUInt16BigEndian(initData.AsSpan(0, 2), 0x0001);
         BinaryPrimitives.WriteUInt16BigEndian(initData.AsSpan(2, 2), HostMaxPacketSize);
-        response = SendSinglePacket(PacketId.Initialization, (ushort)_sequence, PacketFlag.None, initData, initData.Length, _maxTransmissionAttempts, out _);
+        response = SendSinglePacket(PacketId.Initialization, (ushort)_sequence, PacketFlag.None, initData, initData.Length, initAttempts, out _);
         if (response.Length < 4) throw new Exception("Invalid initialization response from target.");
 
         ushort version = BinaryPrimitives.ReadUInt16BigEndian(response.AsSpan(0, 2));
@@ -187,12 +192,31 @@ private byte[] SendSinglePacket(PacketId id, ushort seq, PacketFlag flag, byte[]
     public byte[] Read(int length)
     {
         if (length <= 0) return Array.Empty<byte>();
+        byte[] buffer = new byte[length];
+        int read = ReadInto(buffer, 0, length);
+        if (read < length)
+        {
+            Array.Resize(ref buffer, read);
+        }
+        return buffer;
+    }
+
+    public int ReadInto(byte[] buffer, int offset, int length)
+    {
+        if (length <= 0) return 0;
+        if (offset < 0 || length < 0 || offset + length > buffer.Length)
+        {
+            throw new ArgumentOutOfRangeException(nameof(length));
+        }
+
         byte[] response = SendDataInternal(PacketId.Fastboot, [], 0, _maxTransmissionAttempts);
         if (response.Length > length)
         {
             throw new Exception("UDP protocol error: receive overflow, target sent too much fastboot data.");
         }
-        return response;
+
+        Buffer.BlockCopy(response, 0, buffer, offset, response.Length);
+        return response.Length;
     }
 
     public long Write(byte[] data, int length)
diff --git a/FirmwareKit.Comm.Fastboot/Backend/Usb/UsbDevice.cs b/FirmwareKit.Comm.Fastboot/Backend/Usb/UsbDevice.cs
@@ -2,14 +2,27 @@
 
 namespace FirmwareKit.Comm.Fastboot.Usb;
 
-public abstract class UsbDevice : IFastbootTransport
+public abstract class UsbDevice : IFastbootBufferedTransport
 {
     public string DevicePath { get; set; } = string.Empty;
     public string? SerialNumber { get; set; }
     public ushort VendorId { get; set; }
     public ushort ProductId { get; set; }
     public UsbDeviceType UsbDeviceType { get; set; }
     public abstract byte[] Read(int length);
+    public virtual int ReadInto(byte[] buffer, int offset, int length)
+    {
+        if (length <= 0) return 0;
+        if (offset < 0 || length < 0 || offset + length > buffer.Length)
+        {
+            throw new ArgumentOutOfRangeException(nameof(length));
+        }
+
+        byte[] data = Read(length);
+        if (data.Length == 0) return 0;
+        Buffer.BlockCopy(data, 0, buffer, offset, data.Length);
+        return data.Length;
+    }
     public abstract long Write(byte[] data, int length);
     public abstract int GetSerialNumber();
     public abstract int CreateHandle();
diff --git a/FirmwareKit.Comm.Fastboot/Command/InternalUploadData.cs b/FirmwareKit.Comm.Fastboot/Command/InternalUploadData.cs
@@ -1,5 +1,7 @@
 
 
+using System.Buffers;
+
 namespace FirmwareKit.Comm.Fastboot;
 
 public partial class FastbootDriver
@@ -15,14 +17,42 @@ public FastbootResponse UploadData(string command, Stream output)
 
         long size = response.DataSize;
         long bytesDownloaded = 0;
-        while (bytesDownloaded < size)
+        byte[] buffer = ArrayPool<byte>.Shared.Rent(OnceSendDataSize);
+        try
+        {
+            while (bytesDownloaded < size)
+            {
+                int toRead = (int)Math.Min(OnceSendDataSize, size - bytesDownloaded);
+                int readLen;
+
+                if (Transport is IFastbootBufferedTransport buffered)
+                {
+                    readLen = buffered.ReadInto(buffer, 0, toRead);
+                }
+                else
+                {
+                    byte[] data = Transport.Read(toRead);
+                    if (data != null && data.Length > 0)
+                    {
+                        readLen = data.Length;
+                        Buffer.BlockCopy(data, 0, buffer, 0, readLen);
+                    }
+                    else
+                    {
+                        readLen = 0;
+                    }
+                }
+
+                if (readLen <= 0) throw new Exception("Unexpected EOF from USB.");
+
+                output.Write(buffer, 0, readLen);
+                bytesDownloaded += readLen;
+                NotifyProgress(bytesDownloaded, size);
+            }
+        }
+        finally
         {
-            int toRead = (int)Math.Min(OnceSendDataSize, size - bytesDownloaded);
-            byte[] data = Transport.Read(toRead);
-            if (data == null || data.Length == 0) throw new Exception("Unexpected EOF from USB.");
-            output.Write(data, 0, data.Length);
-            bytesDownloaded += data.Length;
-            NotifyProgress(bytesDownloaded, size);
+            ArrayPool<byte>.Shared.Return(buffer);
         }
 
         return HandleResponse();
