Default to mmap-based I/O on Windows only

src/NuGet.Core/NuGet.Packaging/PackageExtraction/StreamExtensions.cs

@@ -5,91 +5,137 @@
 using System.Buffers;
 using System.IO;
 using System.IO.MemoryMappedFiles;
+using NuGet.Common;
 
 namespace NuGet.Packaging
 {
     public static class StreamExtensions
     {
-        /**
-        Only files smaller than this value will be mmap'ed
-        */
-        private const long MAX_MMAP_SIZE = 10 * 1024 * 1024;
-
         public static string CopyToFile(this Stream inputStream, string fileFullPath)
         {
-            if (Path.GetFileName(fileFullPath).Length == 0)
+            return Testable.Default.CopyToFile(inputStream, fileFullPath);
+        }
+
+        private static void CopyTo(Stream inputStream, Stream outputStream)
+        {
+            // .NET Framework allocates an unavoidable byte[] when using
+            // Stream.CopyTo. Reimplement it, pulling from the pool similar
+            // to .NET 5.
+
+#if NETFRAMEWORK || NETSTANDARD2_0
+            const int bufferSize = 81920; // Same as Stream.CopyTo
+            byte[] buffer = ArrayPool<byte>.Shared.Rent(bufferSize);
+
+            int bytesRead;
+            while ((bytesRead = inputStream.Read(buffer, offset: 0, buffer.Length)) != 0)
             {
-                Directory.CreateDirectory(fileFullPath);
-                return fileFullPath;
+                outputStream.Write(buffer, offset: 0, bytesRead);
             }
 
-            var directory = Path.GetDirectoryName(fileFullPath);
-            if (!Directory.Exists(directory))
+            ArrayPool<byte>.Shared.Return(buffer);
+#else
+            inputStream.CopyTo(outputStream);
+#endif
+        }
+
+        internal class Testable
+        {
+            // Only files smaller than this value will be mmap'ed
+            private const long MAX_MMAP_SIZE = 10 * 1024 * 1024;
+
+            // Mmap can improve file writing performance, but it can make it slower too.
+            // It all depends on a particular hardware configuration, operating system or anti-virus software.
+            // From our benchmarks we concluded that mmap is a good choice for Windows,
+            // but it is not so for other systems.
+            //
+            // 1 - always use memory-mapped files
+            // 0 - never use memory-mapped files
+            // default - use memory-mapped files on Windows, but not on other systems
+            private const string MMAP_VARIABLE_NAME = "NUGET_PACKAGE_EXTRACTION_USE_MMAP";
+
+            private bool _isMMapEnabled { get; }
+
+            internal Testable(IEnvironmentVariableReader environmentVariableReader)
             {
-                Directory.CreateDirectory(directory);
+                _isMMapEnabled = environmentVariableReader.GetEnvironmentVariable(MMAP_VARIABLE_NAME) switch
+                {
+                    "0" => false,
+                    "1" => true,
+                    _ => RuntimeEnvironmentHelper.IsWindows
+                };
             }
 
-            if (File.Exists(fileFullPath))
+            public static Testable Default { get; } = new Testable(EnvironmentVariableWrapper.Instance);
+
+            public string CopyToFile(Stream inputStream, string fileFullPath)
             {
-                // Log and skip adding file
+                if (Path.GetFileName(fileFullPath).Length == 0)
+                {
+                    Directory.CreateDirectory(fileFullPath);
+                    return fileFullPath;
+                }
+
+                var directory = Path.GetDirectoryName(fileFullPath);
+                if (!Directory.Exists(directory))
+                {
+                    Directory.CreateDirectory(directory);
+                }
+
+                if (File.Exists(fileFullPath))
+                {
+                    // Log and skip adding file
+                    return fileFullPath;
+                }
+
+                // For files of a certain size, we can do some Cleverness and mmap
+                // them instead of writing directly to disk. This can improve
+                // performance by a lot on some operating systems and hardware,
+                // particularly Windows
+                long? size = null;
+                try
+                {
+                    size = inputStream.Length;
+                }
+                catch (NotSupportedException)
+                {
+                    // If we can't get Length, just move on.
+                }
+
+                if (_isMMapEnabled && size > 0 && size <= MAX_MMAP_SIZE)
+                {
+                    MmapCopy(inputStream, fileFullPath, size.Value);
+                }
+                else
+                {
+                    FileStreamCopy(inputStream, fileFullPath);
+                }
+
                 return fileFullPath;
             }
 
-            // For files of a certain size, we can do some Cleverness and mmap
-            // them instead of writing directly to disk. This can improve
-            // performance by a lot on some operating systems and hardware,
-            // particularly Windows
-            long? size = null;
-            try
-            {
-                size = inputStream.Length;
-            }
-            catch (NotSupportedException)
-            {
-                // If we can't get Length, just move on.
-            }
-            using (var outputStream = NuGetExtractionFileIO.CreateFile(fileFullPath))
+            internal virtual void MmapCopy(Stream inputStream, string fileFullPath, long size)
             {
-                if (size > 0 && size <= MAX_MMAP_SIZE)
+                using (var outputStream = NuGetExtractionFileIO.CreateFile(fileFullPath))
                 {
                     // NOTE: Linux can't create a mmf from outputStream, so we
                     // need to close the file (which now has the desired
                     // perms), and then re-open it as a memory-mapped file.
                     outputStream.Dispose();
-                    using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(fileFullPath, FileMode.Open, mapName: null, (long)size))
+                    using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(fileFullPath, FileMode.Open, mapName: null, size))
                     using (MemoryMappedViewStream mmstream = mmf.CreateViewStream())
                     {
                         CopyTo(inputStream, mmstream);
                     }
                 }
-                else
-                {
-                    CopyTo(inputStream, outputStream);
-                }
             }
-            return fileFullPath;
-        }
 
-        private static void CopyTo(Stream inputStream, Stream outputStream)
-        {
-            // .NET Framework allocates an unavoidable byte[] when using
-            // Stream.CopyTo. Reimplement it, pulling from the pool similar
-            // to .NET 5.
-
-#if NETFRAMEWORK || NETSTANDARD2_0
-            const int bufferSize = 81920; // Same as Stream.CopyTo
-            byte[] buffer = ArrayPool<byte>.Shared.Rent(bufferSize);
-
-            int bytesRead;
-            while ((bytesRead = inputStream.Read(buffer, offset: 0, buffer.Length)) != 0)
+            internal virtual void FileStreamCopy(Stream inputStream, string fileFullPath)
             {
-                outputStream.Write(buffer, offset: 0, bytesRead);
+                using (var outputStream = NuGetExtractionFileIO.CreateFile(fileFullPath))
+                {
+                    CopyTo(inputStream, outputStream);
+                }
             }
-
-            ArrayPool<byte>.Shared.Return(buffer);
-#else
-            inputStream.CopyTo(outputStream);
-#endif
         }
     }
 }

test/NuGet.Core.Tests/NuGet.Packaging.Test/PackageExtraction/StreamExtensionsTests.cs

@@ -0,0 +1,87 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.IO;
+using System.Text;
+using Moq;
+using NuGet.Common;
+using NuGet.Test.Utility;
+using Xunit;
+
+namespace NuGet.Packaging.Test.PackageExtraction
+{
+    public class StreamExtensionsTests
+    {
+        private const string TestText = "Hello world";
+
+        class TestableProxy : NuGet.Packaging.StreamExtensions.Testable
+        {
+            public bool MmapCopyWasCalled { get; set; }
+            public bool FileStreamCopyWasCalled { get; set; }
+            internal TestableProxy(IEnvironmentVariableReader environmentVariableReader) : base(environmentVariableReader)
+            {
+            }
+
+            internal override void MmapCopy(Stream inputStream, string fileFullPath, long size)
+            {
+                MmapCopyWasCalled = true;
+                base.MmapCopy(inputStream, fileFullPath, size);
+            }
+
+            internal override void FileStreamCopy(Stream inputStream, string fileFullPath)
+            {
+                FileStreamCopyWasCalled = true;
+                base.FileStreamCopy(inputStream, fileFullPath);
+            }
+        }
+
+        [PlatformFact(Platform.Windows)]
+        public void CopyToFile_Windows_CallsMmapCopy()
+        {
+            using (var directory = TestDirectory.Create())
+            {
+                var testPath = Path.Combine(directory, Path.GetRandomFileName());
+                var environmentVariableReader = new Mock<IEnvironmentVariableReader>();
+                var uut = new TestableProxy(environmentVariableReader.Object);
+                uut.CopyToFile(new MemoryStream(Encoding.UTF8.GetBytes(TestText)), testPath);
+                Assert.True(uut.MmapCopyWasCalled);
+                Assert.False(uut.FileStreamCopyWasCalled);
+                Assert.Equal(TestText, File.ReadAllText(testPath));
+            }
+        }
+
+        [PlatformFact(SkipPlatform = Platform.Windows)]
+        public void CopyToFile_NonWindows_CallsFileStreamCopy()
+        {
+            using (var directory = TestDirectory.Create())
+            {
+                var testPath = Path.Combine(directory, Path.GetRandomFileName());
+                var environmentVariableReader = new Mock<IEnvironmentVariableReader>();
+                var uut = new TestableProxy(environmentVariableReader.Object);
+                uut.CopyToFile(new MemoryStream(Encoding.UTF8.GetBytes(TestText)), testPath);
+                Assert.False(uut.MmapCopyWasCalled);
+                Assert.True(uut.FileStreamCopyWasCalled);
+                Assert.Equal(TestText, File.ReadAllText(testPath));
+            }
+        }
+
+        [Theory]
+        [InlineData("0", false)]
+        [InlineData("1", true)]
+        public void CopyToFile_EnvironmentVariable_IsRespected(string env, bool expectedMMap)
+        {
+            using (var directory = TestDirectory.Create())
+            {
+                var environmentVariableReader = new Mock<IEnvironmentVariableReader>();
+                environmentVariableReader.Setup(x => x.GetEnvironmentVariable("NUGET_PACKAGE_EXTRACTION_USE_MMAP"))
+                    .Returns(env);
+                var uut = new TestableProxy(environmentVariableReader.Object);
+                var testPath = Path.Combine(directory, Path.GetRandomFileName());
+                uut.CopyToFile(new MemoryStream(Encoding.UTF8.GetBytes(TestText)), testPath);
+                Assert.Equal(uut.MmapCopyWasCalled, expectedMMap);
+                Assert.Equal(uut.FileStreamCopyWasCalled, !uut.MmapCopyWasCalled);
+                Assert.Equal(TestText, File.ReadAllText(testPath));
+            }
+        }
+    }
+}