[release/9.0-staging] Fix few RandomAccess.Write edge case bugs (#109646)

Co-authored-by: NicoAvanzDev <[email protected]>
Co-authored-by: Adeel Mujahid <[email protected]>
Co-authored-by: Stephen Toub <[email protected]>
Co-authored-by: Michał Petryka <[email protected]>

Author: 5 people

src/libraries/System.Private.CoreLib/src/System/IO/RandomAccess.Unix.cs

@@ -168,37 +168,30 @@ internal static unsafe void WriteGatherAtOffset(SafeFileHandle handle, IReadOnly
 
             var handles = new MemoryHandle[buffersCount];
             Span<Interop.Sys.IOVector> vectors = buffersCount <= IovStackThreshold ?
-                stackalloc Interop.Sys.IOVector[IovStackThreshold] :
+                stackalloc Interop.Sys.IOVector[IovStackThreshold].Slice(0, buffersCount) :
                 new Interop.Sys.IOVector[buffersCount];
 
             try
             {
-                int buffersOffset = 0, firstBufferOffset = 0;
-                while (true)
+                long totalBytesToWrite = 0;
+                for (int i = 0; i < buffersCount; i++)
                 {
-                    long totalBytesToWrite = 0;
-
-                    for (int i = buffersOffset; i < buffersCount; i++)
-                    {
-                        ReadOnlyMemory<byte> buffer = buffers[i];
-                        totalBytesToWrite += buffer.Length;
-
-                        MemoryHandle memoryHandle = buffer.Pin();
-                        vectors[i] = new Interop.Sys.IOVector { Base = firstBufferOffset + (byte*)memoryHandle.Pointer, Count = (UIntPtr)(buffer.Length - firstBufferOffset) };
-                        handles[i] = memoryHandle;
-
-                        firstBufferOffset = 0;
-                    }
+                    ReadOnlyMemory<byte> buffer = buffers[i];
+                    totalBytesToWrite += buffer.Length;
 
-                    if (totalBytesToWrite == 0)
-                    {
-                        break;
-                    }
+                    MemoryHandle memoryHandle = buffer.Pin();
+                    vectors[i] = new Interop.Sys.IOVector { Base = (byte*)memoryHandle.Pointer, Count = (UIntPtr)buffer.Length };
+                    handles[i] = memoryHandle;
+                }
 
+                int buffersOffset = 0;
+                while (totalBytesToWrite > 0)
+                {
                     long bytesWritten;
-                    fixed (Interop.Sys.IOVector* pinnedVectors = &MemoryMarshal.GetReference(vectors))
+                    Span<Interop.Sys.IOVector> left = vectors.Slice(buffersOffset);
+                    fixed (Interop.Sys.IOVector* pinnedVectors = &MemoryMarshal.GetReference(left))
                     {
-                        bytesWritten = Interop.Sys.PWriteV(handle, pinnedVectors, buffersCount, fileOffset);
+                        bytesWritten = Interop.Sys.PWriteV(handle, pinnedVectors, left.Length, fileOffset);
                     }
 
                     FileStreamHelpers.CheckFileCall(bytesWritten, handle.Path);
@@ -208,22 +201,31 @@ internal static unsafe void WriteGatherAtOffset(SafeFileHandle handle, IReadOnly
                     }
 
                     // The write completed successfully but for fewer bytes than requested.
+                    // We need to perform next write where the previous one has finished.
+                    fileOffset += bytesWritten;
+                    totalBytesToWrite -= bytesWritten;
                     // We need to try again for the remainder.
-                    for (int i = 0; i < buffersCount; i++)
+                    while (buffersOffset < buffersCount && bytesWritten > 0)
                     {
-                        int n = buffers[i].Length;
+                        int n = (int)vectors[buffersOffset].Count;
                         if (n <= bytesWritten)
                         {
-                            buffersOffset++;
                             bytesWritten -= n;
-                            if (bytesWritten == 0)
-                            {
-                                break;
-                            }
+                            buffersOffset++;
                         }
                         else
                         {
-                            firstBufferOffset = (int)(bytesWritten - n);
+                            // A partial read: the vector needs to point to the new offset.
+                            // But that offset needs to be relative to the previous attempt.
+                            // Example: we have a single buffer with 30 bytes and the first read returned 10.
+                            // The next read should try to read the remaining 20 bytes, but in case it also reads just 10,
+                            // the third attempt should read last 10 bytes (not 20 again).
+                            Interop.Sys.IOVector current = vectors[buffersOffset];
+                            vectors[buffersOffset] = new Interop.Sys.IOVector
+                            {
+                                Base = current.Base + (int)(bytesWritten),
+                                Count = current.Count - (UIntPtr)(bytesWritten)
+                            };
                             break;
                         }
                     }

src/libraries/System.Private.CoreLib/src/System/IO/RandomAccess.Windows.cs

@@ -435,7 +435,7 @@ internal static long ReadScatterAtOffset(SafeFileHandle handle, IReadOnlyList<Me
         internal static void WriteGatherAtOffset(SafeFileHandle handle, IReadOnlyList<ReadOnlyMemory<byte>> buffers, long fileOffset)
         {
             // WriteFileGather does not support sync handles, so we just call WriteFile in a loop
-            int bytesWritten = 0;
+            long bytesWritten = 0;
             int buffersCount = buffers.Count;
             for (int i = 0; i < buffersCount; i++)
             {

src/libraries/System.Runtime/tests/System.IO.FileSystem.Tests/RandomAccess/WriteGatherAsync.cs

@@ -1,17 +1,20 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Buffers;
 using System.Collections.Generic;
 using System.Linq;
 using System.Security.Cryptography;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.DotNet.XUnitExtensions;
 using Microsoft.Win32.SafeHandles;
 using Xunit;
 
 namespace System.IO.Tests
 {
     [SkipOnPlatform(TestPlatforms.Browser, "async file IO is not supported on browser")]
+    [Collection(nameof(DisableParallelization))] // don't run in parallel, as some of these tests use a LOT of resources
     public class RandomAccess_WriteGatherAsync : RandomAccess_Base<ValueTask>
     {
         protected override ValueTask MethodUnderTest(SafeFileHandle handle, byte[] bytes, long fileOffset)
@@ -133,5 +136,172 @@ public async Task DuplicatedBufferDuplicatesContentAsync(FileOptions options)
             Assert.Equal(repeatCount, actualContent.Length);
             Assert.All(actualContent, actual => Assert.Equal(value, actual));
         }
+
+        [OuterLoop("It consumes a lot of resources (disk space and memory).")]
+        [ConditionalTheory(typeof(PlatformDetection), nameof(PlatformDetection.Is64BitProcess), nameof(PlatformDetection.IsReleaseRuntime))]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, true)]
+        [InlineData(true, false)]
+        public async Task NoInt32OverflowForLargeInputs(bool asyncFile, bool asyncMethod)
+        {
+            // We need to write more than Int32.MaxValue bytes to the disk to reproduce the problem.
+            // To reduce the number of used memory, we allocate only one write buffer and simply repeat it multiple times.
+            // For reading, we need unique buffers to ensure that all of them are getting populated with the right data.
+
+            const int BufferCount = 1002;
+            const int BufferSize = int.MaxValue / 1000;
+            const long FileSize = (long)BufferCount * BufferSize;
+            string filePath = GetTestFilePath();
+
+            FileOptions options = asyncFile ? FileOptions.Asynchronous : FileOptions.None; // we need to test both code paths
+            options |= FileOptions.DeleteOnClose;
+
+            SafeFileHandle? sfh;
+            try
+            {
+                sfh = File.OpenHandle(filePath, FileMode.CreateNew, FileAccess.ReadWrite, FileShare.None, options, preallocationSize: FileSize);
+            }
+            catch (IOException)
+            {
+                throw new SkipTestException("Not enough disk space.");
+            }
+
+            using (sfh)
+            {
+                ReadOnlyMemory<byte> writeBuffer = RandomNumberGenerator.GetBytes(BufferSize);
+                List<ReadOnlyMemory<byte>> writeBuffers = Enumerable.Repeat(writeBuffer, BufferCount).ToList();
+
+                List<NativeMemoryManager> memoryManagers = new List<NativeMemoryManager>(BufferCount);
+                List<Memory<byte>> readBuffers = new List<Memory<byte>>(BufferCount);
+
+                try
+                {
+                    try
+                    {
+                        for (int i = 0; i < BufferCount; i++)
+                        {
+                            // We are using native memory here to get OOM as soon as possible.
+                            NativeMemoryManager nativeMemoryManager = new(BufferSize);
+                            memoryManagers.Add(nativeMemoryManager);
+                            readBuffers.Add(nativeMemoryManager.Memory);
+                        }
+                    }
+                    catch (OutOfMemoryException)
+                    {
+                        throw new SkipTestException("Not enough memory.");
+                    }
+
+                    await Verify(asyncMethod, FileSize, sfh, writeBuffer, writeBuffers, readBuffers);
+                }
+                finally
+                {
+                    foreach (IDisposable memoryManager in memoryManagers)
+                    {
+                        memoryManager.Dispose();
+                    }
+                }
+            }
+
+            static async Task Verify(bool asyncMethod, long FileSize, SafeFileHandle sfh, ReadOnlyMemory<byte> writeBuffer, List<ReadOnlyMemory<byte>> writeBuffers, List<Memory<byte>> readBuffers)
+            {
+                if (asyncMethod)
+                {
+                    await RandomAccess.WriteAsync(sfh, writeBuffers, 0);
+                }
+                else
+                {
+                    RandomAccess.Write(sfh, writeBuffers, 0);
+                }
+
+                Assert.Equal(FileSize, RandomAccess.GetLength(sfh));
+
+                long fileOffset = 0;
+                while (fileOffset < FileSize)
+                {
+                    long bytesRead = asyncMethod
+                        ? await RandomAccess.ReadAsync(sfh, readBuffers, fileOffset)
+                        : RandomAccess.Read(sfh, readBuffers, fileOffset);
+
+                    Assert.InRange(bytesRead, 0, FileSize);
+
+                    while (bytesRead > 0)
+                    {
+                        Memory<byte> readBuffer = readBuffers[0];
+                        if (bytesRead >= readBuffer.Length)
+                        {
+                            AssertExtensions.SequenceEqual(writeBuffer.Span, readBuffer.Span);
+
+                            bytesRead -= readBuffer.Length;
+                            fileOffset += readBuffer.Length;
+
+                            readBuffers.RemoveAt(0);
+                        }
+                        else
+                        {
+                            // A read has finished somewhere in the middle of one of the read buffers.
+                            // Example: buffer had 30 bytes and only 10 were read.
+                            // We don't read the missing part, but try to read the whole buffer again.
+                            // It's not optimal from performance perspective, but it keeps the test logic simple.
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+
+        [Theory]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, true)]
+        [InlineData(true, false)]
+        public async Task IovLimitsAreRespected(bool asyncFile, bool asyncMethod)
+        {
+            // We need to write and read more than IOV_MAX buffers at a time.
+            // IOV_MAX typical value is 1024.
+            const int BufferCount = 1026;
+            const int BufferSize = 1; // the less resources we use, the better
+            const int FileSize = BufferCount * BufferSize;
+            
+            ReadOnlyMemory<byte> writeBuffer = RandomNumberGenerator.GetBytes(BufferSize);
+            ReadOnlyMemory<byte>[] writeBuffers = Enumerable.Repeat(writeBuffer, BufferCount).ToArray();
+            Memory<byte>[] readBuffers = Enumerable.Range(0, BufferCount).Select(_ => new byte[BufferSize].AsMemory()).ToArray();
+
+            FileOptions options = asyncFile ? FileOptions.Asynchronous : FileOptions.None; // we need to test both code paths
+            options |= FileOptions.DeleteOnClose;
+
+            using SafeFileHandle sfh = File.OpenHandle(GetTestFilePath(), FileMode.CreateNew, FileAccess.ReadWrite, FileShare.None, options);
+
+            if (asyncMethod)
+            {
+                await RandomAccess.WriteAsync(sfh, writeBuffers, 0);
+            }
+            else
+            {
+                RandomAccess.Write(sfh, writeBuffers, 0);
+            }
+
+            Assert.Equal(FileSize, RandomAccess.GetLength(sfh));
+
+            long fileOffset = 0;
+            int bufferOffset = 0;
+            while (fileOffset < FileSize)
+            {
+                ArraySegment<Memory<byte>> left = new ArraySegment<Memory<byte>>(readBuffers, bufferOffset, readBuffers.Length - bufferOffset);
+
+                long bytesRead = asyncMethod
+                    ? await RandomAccess.ReadAsync(sfh, left, fileOffset)
+                    : RandomAccess.Read(sfh, left, fileOffset);
+
+                fileOffset += bytesRead;
+                // The following operation is correct only because the BufferSize is 1.
+                bufferOffset += (int)bytesRead;
+            }
+
+            for (int i = 0; i < BufferCount; ++i)
+            {
+                Assert.Equal(writeBuffers[i], readBuffers[i]);
+            }
+        }
     }
 }

src/native/libs/System.Native/pal_io.c

@@ -1883,6 +1883,53 @@ int32_t SystemNative_PWrite(intptr_t fd, void* buffer, int32_t bufferSize, int64
     return (int32_t)count;
 }
 
+#if (HAVE_PREADV || HAVE_PWRITEV) && !defined(TARGET_WASM)
+static int GetAllowedVectorCount(IOVector* vectors, int32_t vectorCount)
+{
+#if defined(IOV_MAX)
+    const int IovMax = IOV_MAX;
+#else
+    // In theory all the platforms that we support define IOV_MAX,
+    // but we want to be extra safe and provde a fallback
+    // in case it turns out to not be true.
+    // 16 is low, but supported on every platform.
+    const int IovMax = 16;
+#endif
+
+    int allowedCount = (int)vectorCount;
+
+    // We need to respect the limit of items that can be passed in iov.
+    // In case of writes, the managed code is responsible for handling incomplete writes.
+    // In case of reads, we simply returns the number of bytes read and it's up to the users.
+    if (IovMax < allowedCount)
+    {
+        allowedCount = IovMax;
+    }
+
+#if defined(TARGET_APPLE)
+    // For macOS preadv and pwritev can fail with EINVAL when the total length
+    // of all vectors overflows a 32-bit integer.
+    size_t totalLength = 0;
+    for (int i = 0; i < allowedCount; i++) 
+    {
+        assert(INT_MAX >= vectors[i].Count);
+
+        totalLength += vectors[i].Count;
+
+        if (totalLength > INT_MAX)
+        {
+            allowedCount = i;
+            break;
+        }
+    }
+#else
+    (void)vectors;
+#endif
+
+    return allowedCount;
+}
+#endif // (HAVE_PREADV || HAVE_PWRITEV) && !defined(TARGET_WASM)
+
 int64_t SystemNative_PReadV(intptr_t fd, IOVector* vectors, int32_t vectorCount, int64_t fileOffset)
 {
     assert(vectors != NULL);
@@ -1891,7 +1938,8 @@ int64_t SystemNative_PReadV(intptr_t fd, IOVector* vectors, int32_t vectorCount,
     int64_t count = 0;
     int fileDescriptor = ToFileDescriptor(fd);
 #if HAVE_PREADV && !defined(TARGET_WASM) // preadv is buggy on WASM
-    while ((count = preadv(fileDescriptor, (struct iovec*)vectors, (int)vectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
+    int allowedVectorCount = GetAllowedVectorCount(vectors, vectorCount);
+    while ((count = preadv(fileDescriptor, (struct iovec*)vectors, allowedVectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
 #else
     int64_t current;
     for (int i = 0; i < vectorCount; i++)
@@ -1931,7 +1979,8 @@ int64_t SystemNative_PWriteV(intptr_t fd, IOVector* vectors, int32_t vectorCount
     int64_t count = 0;
     int fileDescriptor = ToFileDescriptor(fd);
 #if HAVE_PWRITEV && !defined(TARGET_WASM) // pwritev is buggy on WASM
-    while ((count = pwritev(fileDescriptor, (struct iovec*)vectors, (int)vectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
+    int allowedVectorCount = GetAllowedVectorCount(vectors, vectorCount);
+    while ((count = pwritev(fileDescriptor, (struct iovec*)vectors, allowedVectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
 #else
     int64_t current;
     for (int i = 0; i < vectorCount; i++)