✨ Move image upload memory allocation to the service.

I honestly don't like it, but let's hope this is enough for things to work.

src/Exo/Service/Exo.Service.Core/ImageStorageService.cs

@@ -101,6 +101,16 @@ public async IAsyncEnumerable<ImageChangeNotification> WatchChangesAsync([Enumer
 		}
 	}
 
+	public async ValueTask<bool> HasImageAsync(string imageName, CancellationToken cancellationToken)
+	{
+		if (!ImageNameSerializer.IsNameValid(imageName)) throw new ArgumentException("Invalid name.");
+
+		using (await _lock.WaitAsync(cancellationToken).ConfigureAwait(false))
+		{
+			return _imageCollection.ContainsKey(imageName);
+		}
+	}
+
 	public async ValueTask AddImageAsync(string imageName, ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
 	{
 		if (!ImageNameSerializer.IsNameValid(imageName)) throw new ArgumentException("Invalid name.");

src/Exo/Service/Exo.Service.Grpc/GrpcImageService.cs

@@ -8,6 +8,7 @@ namespace Exo.Service.Grpc;
 internal sealed class GrpcImageService : IImageService
 {
 	private readonly ImageStorageService _imageStorageService;
+	private ImageAddState? _imageAddState;
 
 	public GrpcImageService(ImageStorageService imagesService) => _imageStorageService = imagesService;
 
@@ -23,17 +24,100 @@ internal sealed class GrpcImageService : IImageService
 		}
 	}
 
-	public async ValueTask AddImageAsync(ImageRegistrationRequest request, CancellationToken cancellationToken)
+	public async IAsyncEnumerable<ImageRegistrationBeginResponse> BeginAddImageAsync(ImageRegistrationBeginRequest request, [EnumeratorCancellation] CancellationToken cancellationToken)
 	{
-		using (var memoryMappedFile = MemoryMappedFile.CreateOrOpen(request.SharedMemoryName, (long)request.SharedMemoryLength, MemoryMappedFileAccess.Read))
-		using (var memoryManager = new MemoryMappedFileMemoryManager(memoryMappedFile, 0, (int)request.SharedMemoryLength, MemoryMappedFileAccess.Read))
+		if (await _imageStorageService.HasImageAsync(request.ImageName, cancellationToken).ConfigureAwait(false))
 		{
-			await _imageStorageService.AddImageAsync(request.ImageName, memoryManager.Memory, cancellationToken).ConfigureAwait(false);
+			throw new InvalidOperationException("An image with the specified name is already present.");
 		}
+		// 24MB is the maximum image size that could be theoretically allocated on the NZXT Kraken Z, so it seems reasonable to fix this as the limit for now?
+		// This number is semi-arbitrary, as we have to allow somewhat large images, but we don't want to allow anything crazy.
+		// The risk is someone abusing the service to waste memory. Stupid but possible.
+		// We have to be careful because apps cannot create shared memory in the Global namespace themselves, so the service needs to handle it ☹️
+		// This is also the reason why only a single parallel request is allowed. Anyway, this would never be a problem in practice as the UI is supposed to be the unique client.
+		if (request.Length > 24 * 1024 * 1024) throw new Exception("Maximum allowed image size exceeded.");
+
+		var state = new ImageAddState(request.ImageName);
+
+		if (Interlocked.CompareExchange(ref _imageAddState, state, null) is not null)
+		{
+			state.Dispose();
+			throw new InvalidOperationException("Another operation is currently pending.");
+		}
+
+		try
+		{
+			yield return new() { RequestId = state.RequestId, SharedMemoryName = state.Initialize(request.Length) };
+			await state.WaitForWriteCompletion().WaitAsync(cancellationToken).ConfigureAwait(false);
+			try
+			{
+				using (var memoryManager = state.CreateMemoryManagerForRead())
+				{
+					await _imageStorageService.AddImageAsync(request.ImageName, memoryManager.Memory, cancellationToken).ConfigureAwait(false);
+				}
+				state.TryNotifyReadCompletion();
+			}
+			catch (Exception ex)
+			{
+				state.TryNotifyReadException(ex);
+			}
+		}
+		finally
+		{
+			state.Dispose();
+			Interlocked.CompareExchange(ref _imageAddState, null, state);
+		}
+	}
+
+	public async ValueTask EndAddImageAsync(ImageRegistrationEndRequest request, CancellationToken cancellationToken)
+	{
+		if (Volatile.Read(ref _imageAddState) is not { } state || state.RequestId != request.RequestId || !state.TryNotifyWriteCompletion()) throw new InvalidOperationException();
+
+		await state.WaitForReadCompletion().WaitAsync(cancellationToken).ConfigureAwait(false);
 	}
 
 	public async ValueTask RemoveImageAsync(ImageReference request, CancellationToken cancellationToken)
 	{
 		await _imageStorageService.RemoveImageAsync(request.ImageName, cancellationToken).ConfigureAwait(false);
 	}
+
+	private sealed class ImageAddState : IDisposable
+	{
+		private readonly Guid _requestId;
+		private readonly string _imageName;
+		private TaskCompletionSource<bool>? _writeTaskCompletionSource;
+		private TaskCompletionSource<bool>? _readTaskCompletionSource;
+		private SharedMemory? _sharedMemory;
+
+		public ImageAddState(string imageName)
+		{
+			_requestId = Guid.NewGuid();
+			_imageName = imageName;
+		}
+
+		public void Dispose()
+		{
+			if (Interlocked.Exchange(ref _writeTaskCompletionSource, null) is { } tcs1 && !tcs1.Task.IsCompleted) tcs1.TrySetResult(false);
+			if (Interlocked.Exchange(ref _readTaskCompletionSource, null) is { } tcs2 && !tcs2.Task.IsCompleted) tcs2.TrySetResult(false);
+			if (Interlocked.Exchange(ref _sharedMemory, null) is { } sharedMemory) sharedMemory.Dispose();
+		}
+
+		public Guid RequestId => _requestId;
+
+		public MemoryMappedFileMemoryManager CreateMemoryManagerForRead() => (_sharedMemory ?? throw new InvalidOperationException()).CreateMemoryManager(MemoryMappedFileAccess.Read);
+
+		public string Initialize(uint length)
+		{
+			_writeTaskCompletionSource = new();
+			_sharedMemory = SharedMemory.Create("Exo_Image_", length);
+			return _sharedMemory.Name;
+		}
+
+		public bool TryNotifyWriteCompletion() => _writeTaskCompletionSource?.TrySetResult(true) ?? false;
+		public bool TryNotifyReadCompletion() => _readTaskCompletionSource?.TrySetResult(true) ?? false;
+		public bool TryNotifyReadException(Exception ex) => _readTaskCompletionSource?.TrySetException(ex) ?? false;
+
+		public Task<bool> WaitForWriteCompletion() => _writeTaskCompletionSource?.Task ?? Task.FromResult(false);
+		public Task<bool> WaitForReadCompletion() => _readTaskCompletionSource?.Task ?? Task.FromResult(true);
+	}
 }

src/Exo/Service/Exo.Service/Startup.cs

@@ -262,7 +262,7 @@ public void ConfigureServices(IServiceCollection services)
 		services.AddSingleton<GrpcServiceLifetimeService>();
 		services.AddSingleton<IOverlayCustomMenuService>(sp => sp.GetRequiredService<GrpcCustomMenuService>());
 		services.AddSingleton<ISettingsCustomMenuService>(sp => sp.GetRequiredService<GrpcCustomMenuService>());
-		services.AddCodeFirstGrpc();
+		services.AddCodeFirstGrpc(options => options.MaxReceiveMessageSize = 512 * 1024);
 	}
 
 	// This method gets called by the runtime. Use this method to configure the HTTP request pipeline.

src/Exo/Ui/Exo.Contracts.Ui.Settings/IImageService.cs

@@ -14,9 +14,23 @@ public interface IImageService
 	[OperationContract(Name = "WatchImages")]
 	IAsyncEnumerable<WatchNotification<ImageInformation>> WatchImagesAsync(CancellationToken cancellationToken);
 
-	// TODO: Maybe surface an error code so that the reason for a fail are more easily accessible.
-	[OperationContract(Name = "AddImage")]
-	ValueTask AddImageAsync(ImageRegistrationRequest request, CancellationToken cancellationToken);
+	/// <summary>Begins an image upload to the service by requesting a buffer for the specified image.</summary>
+	/// <remarks>
+	/// This method will asynchronously allocate a shared memory buffer and keep it open until the call is completed.
+	/// As such, the stream will only return a single value, then the method will wait either cancellation or proper completion through a call to <see cref="EndAddImageAsync(ImageRegistrationEndRequest, CancellationToken)"/>.
+	/// </remarks>
+	/// <param name="request"></param>
+	/// <param name="cancellationToken"></param>
+	/// <returns></returns>
+	[OperationContract(Name = "BeginAddImage")]
+	IAsyncEnumerable<ImageRegistrationBeginResponse> BeginAddImageAsync(ImageRegistrationBeginRequest request, CancellationToken cancellationToken);
+
+	/// <summary>Completes an image upload after having written the data to the shared memory buffer.</summary>
+	/// <param name="response"></param>
+	/// <param name="cancellationToken"></param>
+	/// <returns></returns>
+	[OperationContract(Name = "EndAddImage")]
+	ValueTask EndAddImageAsync(ImageRegistrationEndRequest request, CancellationToken cancellationToken);
 
 	[OperationContract(Name = "RemoveImage")]
 	ValueTask RemoveImageAsync(ImageReference request, CancellationToken cancellationToken);

src/Exo/Ui/Exo.Contracts.Ui.Settings/ImageRegistrationBeginRequest.cs

@@ -0,0 +1,12 @@
+using System.Runtime.Serialization;
+
+namespace Exo.Contracts.Ui.Settings;
+
+[DataContract]
+public sealed class ImageRegistrationBeginRequest
+{
+	[DataMember(Order = 1)]
+	public required string ImageName { get; init; }
+	[DataMember(Order = 2)]
+	public required uint Length { get; init; }
+}

src/Exo/Ui/Exo.Contracts.Ui.Settings/ImageRegistrationRequest.cs → src/Exo/Ui/Exo.Contracts.Ui.Settings/ImageRegistrationBeginResponse.cs

@@ -3,12 +3,10 @@
 namespace Exo.Contracts.Ui.Settings;
 
 [DataContract]
-public sealed class ImageRegistrationRequest
+public sealed class ImageRegistrationBeginResponse
 {
 	[DataMember(Order = 1)]
-	public required string ImageName { get; init; }
+	public required Guid RequestId { get; init; }
 	[DataMember(Order = 2)]
 	public required string SharedMemoryName { get; init; }
-	[DataMember(Order = 3)]
-	public required ulong SharedMemoryLength { get; init; }
 }

src/Exo/Ui/Exo.Contracts.Ui.Settings/ImageRegistrationEndRequest.cs

@@ -0,0 +1,10 @@
+using System.Runtime.Serialization;
+
+namespace Exo.Contracts.Ui.Settings;
+
+[DataContract]
+public sealed class ImageRegistrationEndRequest
+{
+	[DataMember(Order = 1)]
+	public required Guid RequestId { get; init; }
+}

src/Exo/Ui/Exo.Settings.Ui/ImagesPage.xaml

@@ -22,7 +22,7 @@
                     <ColumnDefinition Width="*" MaxWidth="200" />
                     <ColumnDefinition Width="Auto" />
                 </Grid.ColumnDefinitions>
-                <Image Margin="{ThemeResource RowContentMargin}" Source="{Binding LoadedImageData, Converter={StaticResource SharedMemoryToBitmapImageConverter}}" />
+                <Image Margin="{ThemeResource RowContentMargin}" Source="{Binding LoadedImageData, Converter={StaticResource ByteArrayToBitmapImageConverter}}" />
                 <StackPanel Grid.Column="1" Orientation="Vertical" Margin="{ThemeResource RowContentMargin}">
                     <TextBox Text="{Binding LoadedImageName, Mode=TwoWay}" Header="Name" Width="250" IsEnabled="{Binding LoadedImageData, Converter={StaticResource NullabilityToBooleanConverter}}" />
                     <Button Margin="{ThemeResource RowLabelMargin}" Command="{Binding OpenImageCommand}" Width="120">Open…</Button>

src/Exo/Ui/Exo.Settings.Ui/ImagesPage.xaml.cs

@@ -1,8 +1,5 @@
-using Exo.Memory;
 using Exo.Settings.Ui.ViewModels;
 using Microsoft.UI.Xaml.Controls;
-using Windows.Storage.Pickers;
-using WinRT.Interop;
 
 namespace Exo.Settings.Ui;
 

src/Exo/Ui/Exo.Settings.Ui/ViewModels/ImagesViewModel.cs

@@ -1,5 +1,6 @@
 using System.Buffers;
 using System.Collections.ObjectModel;
+using System.IO.MemoryMappedFiles;
 using System.Security.Cryptography;
 using System.Windows.Input;
 using Exo.Contracts.Ui.Settings;
@@ -73,7 +74,7 @@ public async void Execute(object? parameter)
 
 	private bool _isReady;
 	private string? _loadedImageName;
-	private SharedMemory? _loadedImageData;
+	private byte[]? _loadedImageData;
 
 	private readonly SettingsServiceConnectionManager _connectionManager;
 	private readonly IFileOpenDialog _fileOpenDialog;
@@ -180,23 +181,22 @@ public string? LoadedImageName
 		}
 	}
 
-	public SharedMemory? LoadedImageData
+	public byte[]? LoadedImageData
 	{
 		get => _loadedImageData;
 		private set
 		{
 			if (value != _loadedImageData)
 			{
 				bool couldAddImage = CanAddImage;
-				_loadedImageData?.Dispose();
 				_loadedImageData = value;
 				NotifyPropertyChanged(ChangedProperty.LoadedImageData);
 				if (couldAddImage != CanAddImage) _addImageCommand.NotifyCanExecuteChanged();
 			}
 		}
 	}
 
-	private void SetImage(string name, SharedMemory data)
+	private void SetImage(string name, byte[] data)
 	{
 		LoadedImageName = name;
 		LoadedImageData = data;
@@ -215,7 +215,7 @@ private async Task OpenImageAsync(CancellationToken cancellationToken)
 		var file = await _fileOpenDialog.OpenAsync([".bmp", ".gif", ".png", ".jpg", ".webp",]);
 
 		if (file is null) return;
-		SharedMemory? data;
+		byte[]? data;
 		using (var stream = await file.OpenForReadAsync())
 		{
 			long length = stream.Length;
@@ -225,11 +225,8 @@ private async Task OpenImageAsync(CancellationToken cancellationToken)
 			}
 			else
 			{
-				data = SharedMemory.Create("Exo_Image_", (ulong)length);
-				using (var viewStream = data.CreateWriteStream())
-				{
-					await stream.CopyToAsync(viewStream);
-				}
+				data = new byte[length];
+				await stream.ReadExactlyAsync(data, cancellationToken);
 			}
 		}
 
@@ -253,14 +250,31 @@ private async Task OpenImageAsync(CancellationToken cancellationToken)
 	private async Task AddImageAsync(CancellationToken cancellationToken)
 	{
 		if (_imageService is null || _loadedImageName is null || !IsNameValid(_loadedImageName) || _loadedImageData is null) return;
-
-		await _imageService.AddImageAsync(new() { ImageName = _loadedImageName, SharedMemoryName = _loadedImageData.Name, SharedMemoryLength = _loadedImageData.Length }, cancellationToken);
-		_loadedImageName = null;
-		_loadedImageData.Dispose();
-		_loadedImageData = null;
-		NotifyPropertyChanged(ChangedProperty.LoadedImageName);
-		NotifyPropertyChanged(ChangedProperty.LoadedImageData);
-		_addImageCommand.NotifyCanExecuteChanged();
+		IsNotBusy = false;
+		try
+		{
+			bool isDone = false;
+			await foreach (var response in _imageService.BeginAddImageAsync(new() { ImageName = _loadedImageName, Length = (uint)_loadedImageData.Length }, cancellationToken))
+			{
+				if (isDone) throw new InvalidOperationException();
+				using (var sharedMemory = SharedMemory.Open(response.SharedMemoryName, (uint)_loadedImageData.Length, MemoryMappedFileAccess.Write))
+				using (var memoryManager = sharedMemory.CreateMemoryManager(MemoryMappedFileAccess.Write))
+				{
+					_loadedImageData.AsSpan().CopyTo(memoryManager.GetSpan());
+				}
+				await _imageService.EndAddImageAsync(new() { RequestId = response.RequestId }, cancellationToken);
+				isDone = true;
+			}
+			_loadedImageName = null;
+			_loadedImageData = null;
+			NotifyPropertyChanged(ChangedProperty.LoadedImageName);
+			NotifyPropertyChanged(ChangedProperty.LoadedImageData);
+			_addImageCommand.NotifyCanExecuteChanged();
+		}
+		finally
+		{
+			IsNotBusy = true;
+		}
 	}
 }