Get shared memory directories (in Unix) from AppSettings (#7239)

* Checking env var for directories where MemoryMappedFiles can be created

* Tests for AllowedDirectories and ValidDirectories

* WIP - fixing tests

* WIP - added missing comment

* Cleanup of tests

* Remove directories and subdirectories when cleaning up previously present directories

* Sleep between creating directory and calling into MemoryMappedFileAccessor to ensure a significant measurable delay between the two creations of the directory

* Addressing comments

Author: Gohar Irfan Chaudhry

src/WebJobs.Script.Grpc/Extensions/RpcSharedMemoryDataExtensions.cs

@@ -13,6 +13,11 @@ internal static class RpcSharedMemoryDataExtensions
     {
         internal static async Task<RpcSharedMemory> ToRpcSharedMemoryAsync(this object value, ILogger logger, string invocationId, ISharedMemoryManager sharedMemoryManager)
         {
+            if (value == null)
+            {
+                return new RpcSharedMemory();
+            }
+
             if (!sharedMemoryManager.IsSupported(value))
             {
                 return null;

src/WebJobs.Script/Workers/Rpc/RpcWorkerConstants.cs

@@ -11,6 +11,9 @@ public static class RpcWorkerConstants
         public const string FunctionWorkerRuntimeVersionSettingName = "FUNCTIONS_WORKER_RUNTIME_VERSION";
         public const string FunctionsWorkerProcessCountSettingName = "FUNCTIONS_WORKER_PROCESS_COUNT";
         public const string FunctionsWorkerSharedMemoryDataTransferEnabledSettingName = "FUNCTIONS_WORKER_SHARED_MEMORY_DATA_TRANSFER_ENABLED";
+        // Comma-separated list of directories where shared memory maps can be created for data transfer between host and worker.
+        // This will override the default directories.
+        public const string FunctionsUnixSharedMemoryDirectories = "FUNCTIONS_UNIX_SHARED_MEMORY_DIRECTORIES";
         public const string DotNetLanguageWorkerName = "dotnet";
         public const string NodeLanguageWorkerName = "node";
         public const string JavaLanguageWorkerName = "java";

src/WebJobs.Script/Workers/SharedMemoryDataTransfer/MemoryMappedFileAccessorUnix.cs

@@ -5,7 +5,9 @@
 using System.Collections.Generic;
 using System.IO;
 using System.IO.MemoryMappedFiles;
+using System.Linq;
 using System.Runtime.InteropServices;
+using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.Azure.WebJobs.Script.Workers.SharedMemoryDataTransfer
@@ -15,15 +17,22 @@ namespace Microsoft.Azure.WebJobs.Script.Workers.SharedMemoryDataTransfer
     /// </summary>
     public class MemoryMappedFileAccessorUnix : MemoryMappedFileAccessor
     {
-        public MemoryMappedFileAccessorUnix(ILogger<MemoryMappedFileAccessor> logger) : base(logger)
+        private IEnvironment _environment;
+
+        public MemoryMappedFileAccessorUnix(ILogger<MemoryMappedFileAccessor> logger, IEnvironment environment) : base(logger)
         {
             ValidatePlatform(new List<OSPlatform>()
             {
                 OSPlatform.Linux,
                 OSPlatform.OSX
             });
+
+            _environment = environment;
+            ValidDirectories = GetValidDirectories();
         }
 
+        internal List<string> ValidDirectories { get; private set; }
+
         public override bool TryCreate(string mapName, long size, out MemoryMappedFile mmf)
         {
             mmf = null;
@@ -127,6 +136,74 @@ public override void Delete(string mapName, MemoryMappedFile mmf)
             }
         }
 
+        /// <summary>
+        /// Checks if a list of directories is specified in AppSettings to create <see cref="MemoryMappedFile"/>.
+        /// If one is specified, returns that list. Otherwise returns the default list.
+        /// </summary>
+        /// <returns>List of paths of directories where <see cref="MemoryMappedFile"/> are allowed to be created.</returns>
+        internal List<string> GetAllowedDirectories()
+        {
+            string envVal = _environment.GetEnvironmentVariable(RpcWorkerConstants.FunctionsUnixSharedMemoryDirectories);
+            if (string.IsNullOrEmpty(envVal))
+            {
+                return SharedMemoryConstants.TempDirs;
+            }
+
+            return envVal.Split(',').ToList();
+        }
+
+        /// <summary>
+        /// From a list of allowed directories where <see cref="MemoryMappedFile"/> can be created, it will return
+        /// a list of those that are valid (i.e. exist, or have been successfully created).
+        /// </summary>
+        /// <returns>List of paths of directories where <see cref="MemoryMappedFile"/> can be created.</returns>
+        internal List<string> GetValidDirectories()
+        {
+            List<string> allowedDirectories = GetAllowedDirectories();
+            List<string> validDirectories = new List<string>();
+
+            foreach (string directory in allowedDirectories)
+            {
+                string path = Path.Combine(directory, SharedMemoryConstants.TempDirSuffix);
+                if (Directory.Exists(path))
+                {
+                    Logger.LogTrace("Found directory for shared memory usage: {Directory}", path);
+                    try
+                    {
+                        // If the directory already exists (maybe from a previous run of the host) then clean it up and start afresh
+                        // The previously created memory maps in that directory are not needed and we need to should clean up the memory
+                        Directory.Delete(path, recursive: true);
+                        Logger.LogTrace("Cleaned up existing directory for shared memory usage: {Directory}", path);
+                    }
+                    catch (Exception exception)
+                    {
+                        Logger.LogWarning(exception, "Cannot delete existing directory for shared memory usage: {Directory}", path);
+                    }
+                }
+
+                try
+                {
+                    DirectoryInfo info = Directory.CreateDirectory(path);
+                    if (info.Exists)
+                    {
+                        validDirectories.Add(path);
+                        Logger.LogTrace("Created directory for shared memory usage: {Directory}", path);
+                    }
+                    else
+                    {
+                        Logger.LogWarning("Directory for shared memory usage does not exist: {Directory}", path);
+                    }
+                }
+                catch (Exception exception)
+                {
+                    Logger.LogWarning(exception, "Cannot create directory for shared memory usage: {Directory}", path);
+                }
+            }
+
+            Logger.LogDebug("Valid directories for shared memory usage: {Directories}", string.Join(",", validDirectories));
+            return validDirectories;
+        }
+
         /// <summary>
         /// Get the path of the file to store a <see cref="MemoryMappedFile"/>.
         /// We first try to mount it in memory-mounted directories (e.g. /dev/shm/).
@@ -136,15 +213,11 @@ public override void Delete(string mapName, MemoryMappedFile mmf)
         /// <see cref="null"/> otherwise.</returns>
         private string GetPath(string mapName)
         {
-            // We escape the mapName to make it a valid file name
-            // Python will use urllib.parse.quote_plus(mapName)
-            string escapedMapName = Uri.EscapeDataString(mapName);
-
             // Check if the file already exists
             string filePath;
-            foreach (string tempDir in SharedMemoryConstants.TempDirs)
+            foreach (string tempDir in ValidDirectories)
             {
-                filePath = Path.Combine(tempDir, SharedMemoryConstants.TempDirSuffix, escapedMapName);
+                filePath = Path.Combine(tempDir, mapName);
                 if (File.Exists(filePath))
                 {
                     return filePath;
@@ -162,14 +235,11 @@ private string GetPath(string mapName)
         /// <returns>Created path.</returns>
         private string CreatePath(string mapName, long size)
         {
-            string escapedMapName = Uri.EscapeDataString(mapName);
-
             // Create a new file
-            string newTempDir = GetDirectory(size);
-            if (newTempDir != null)
+            string tempDir = GetDirectory(size);
+            if (tempDir != null)
             {
-                DirectoryInfo newTempDirInfo = Directory.CreateDirectory(newTempDir);
-                return Path.Combine(newTempDirInfo.FullName, escapedMapName);
+                return Path.Combine(tempDir, mapName);
             }
             else
             {
@@ -187,7 +257,7 @@ private string CreatePath(string mapName, long size)
         /// <see cref="SharedMemoryConstants.TempDirSuffix"/> is used.</returns>
         private string GetDirectory(long size)
         {
-            foreach (string tempDir in SharedMemoryConstants.TempDirs)
+            foreach (string tempDir in ValidDirectories)
             {
                 try
                 {
@@ -197,7 +267,7 @@ private string GetDirectory(long size)
                         long minSize = size + SharedMemoryConstants.TempDirMinSize;
                         if (driveInfo.AvailableFreeSpace > minSize)
                         {
-                            return Path.Combine(tempDir, SharedMemoryConstants.TempDirSuffix);
+                            return tempDir;
                         }
                     }
                 }

src/WebJobs.Script/Workers/SharedMemoryDataTransfer/SharedMemoryConstants.cs

@@ -33,13 +33,17 @@ internal class SharedMemoryConstants
         /// <summary>
         /// Minimum size (in number of bytes) an object must be in order for it to be transferred over shared memory.
         /// If the object is smaller than this, gRPC is used.
+        /// Note: This needs to be consistent among the host and workers.
+        ///       e.g. in the Python worker, it is defined in shared_memory_constants.py
         /// </summary>
         public const long MinObjectBytesForSharedMemoryTransfer = 1024 * 1024; // 1 MB
 
         /// <summary>
         /// Maximum size (in number of bytes) an object can be in order for it to be transferred over shared memory.
         /// This limit is imposed because initializing objects like <see cref="byte[]"/> greater than 2GB is not allowed.
         /// Ref: https://stackoverflow.com/a/3944336/3132415
+        /// Note: This needs to be consistent among the host and workers.
+        ///       e.g. in the Python worker, it is defined in shared_memory_constants.py
         /// </summary>
         public const long MaxObjectBytesForSharedMemoryTransfer = ((long)2 * 1024 * 1024 * 1024) - 1; // 2 GB
 

src/WebJobs.Script/Workers/SharedMemoryDataTransfer/SharedMemoryMap.cs

@@ -195,6 +195,12 @@ public void Dispose()
         public async Task<byte[]> CopyStreamAsync(int offset, int count)
         {
             Stream contentStream = await GetStreamAsync();
+            if (contentStream == null)
+            {
+                _logger.LogError("Cannot get stream for shared memory map: {MapName}", _mapName);
+                return null;
+            }
+
             int contentLength = (int)contentStream.Length;
             if (contentLength == 0)
             {