Using a struct in an IAsyncEnumerable parameter with native AOT doesn't work

[eerhardt]

Using the following code:

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

#pragma warning disable SA1402 // File may only contain a single type
#pragma warning disable SA1649 // File name should match first type name

using System.Diagnostics.CodeAnalysis;
using System.Text.Json;
using System.Text.Json.Serialization;
using Nerdbank.Streams;
using StreamJsonRpc;
using StreamJsonRpc.Reflection;

Console.WriteLine("This test is run by \"dotnet publish -r [RID]-x64\" rather than by executing the program.");

// That said, this "program" can run select scenarios to verify that they work in a Native AOT environment.
// When TUnit fixes https://github.com/thomhurst/TUnit/issues/2458, we can move this part of the program to unit tests.
(Stream clientPipe, Stream serverPipe) = FullDuplexStream.CreatePair();
JsonRpc serverRpc = new JsonRpc(new HeaderDelimitedMessageHandler(serverPipe, CreateFormatter()));
JsonRpc clientRpc = new JsonRpc(new HeaderDelimitedMessageHandler(clientPipe, CreateFormatter()));
serverRpc.AddLocalRpcMethod("Add", new Server().Add);
serverRpc.AddLocalRpcMethod("GetOutputsAsync", new Server().GetOutputsAsync);
serverRpc.StartListening();
clientRpc.StartListening();

int sum = await clientRpc.InvokeAsync<int>(nameof(Server.Add), 2, 5);
Console.WriteLine($"2 + 5 = {sum}");

IAsyncEnumerable<CommandOutput> output = await clientRpc.InvokeAsync<IAsyncEnumerable<CommandOutput>>(nameof(Server.GetOutputsAsync));
await foreach (CommandOutput item in output)
{
    Console.WriteLine(item.Text);
}

// When properly configured, this formatter is safe in Native AOT scenarios for
// the very limited use case shown in this program.
[UnconditionalSuppressMessage("Trimming", "IL2026", Justification = "Using the Json source generator.")]
[UnconditionalSuppressMessage("AOT", "IL3050", Justification = "Using the Json source generator.")]
IJsonRpcMessageFormatter CreateFormatter() => new SystemTextJsonFormatter()
{
    JsonSerializerOptions = { TypeInfoResolver = SourceGenerationContext.Default },
};

internal struct CommandOutput
{
    public required string Text { get; init; }
}

internal class Server
{
    public int Add(int a, int b) => a + b;

    public async IAsyncEnumerable<CommandOutput> GetOutputsAsync()
    {
        yield return new CommandOutput { Text = "Output 1" };
        await Task.Delay(1000); // Simulate some delay.
        yield return new CommandOutput { Text = "Output 2" };
        yield return new CommandOutput { Text = "Output 3" };
    }
}

[JsonSerializable(typeof(int))]
[JsonSerializable(typeof(long))]
[JsonSerializable(typeof(JsonElement))]
[JsonSerializable(typeof(IAsyncEnumerable<CommandOutput>))]
[JsonSerializable(typeof(MessageFormatterEnumerableTracker.EnumeratorResults<CommandOutput>))]
internal partial class SourceGenerationContext : JsonSerializerContext;

This application runs correctly with dotnet run.

However, when I dotnet publish it with PublishAot=true, running the same application produces an error:

This test is run by "dotnet publish -r [RID]-x64" rather than by executing the program.
2 + 5 = 7
Unhandled Exception: StreamJsonRpc.ConnectionLostException: The JSON-RPC connection with the remote party was lost before the request could complete.
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw() + 0x20
   at System.Runtime.CompilerServices.TaskAwaiter.ThrowForNonSuccess(Task) + 0xb2
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task, ConfigureAwaitOptions) + 0x4b
   at StreamJsonRpc.JsonRpc.<InvokeCoreAsync>d__170.MoveNext() + 0x7e2
--- End of stack trace from previous location ---
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw() + 0x20
   at System.Runtime.CompilerServices.TaskAwaiter.ThrowForNonSuccess(Task) + 0xb2
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task, ConfigureAwaitOptions) + 0x4b
   at StreamJsonRpc.JsonRpc.<InvokeCoreAsync>d__159`1.MoveNext() + 0x485
--- End of stack trace from previous location ---
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw() + 0x20
   at System.Runtime.CompilerServices.TaskAwaiter.ThrowForNonSuccess(Task) + 0xb2
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task, ConfigureAwaitOptions) + 0x4b
   at Program.<<Main>$>d__0.MoveNext() + 0x469
--- End of stack trace from previous location ---
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw() + 0x20
   at System.Runtime.CompilerServices.TaskAwaiter.ThrowForNonSuccess(Task) + 0xb2
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task, ConfigureAwaitOptions) + 0x4b
   at Program.<Main>(String[] args) + 0x24
   at NativeAOTCompatibility.Test!<BaseAddress>+0x25e9a0

If I change internal struct CommandOutput to internal class CommandOutput it works correctly.

The reason (AFAICT) is because the server is throwing an exception:

System.NotSupportedException: 'StreamJsonRpc.SystemTextJsonFormatter+AsyncEnumerableConverter+Converter`1[CommandOutput]' is missing native code or metadata. This can happen for code that is not compatible with trimming or AOT. Inspect and fix trimming and AOT related warnings that were generated when the app was published. For more information see https://aka.ms/nativeaot-compatibility
   at System.Reflection.Runtime.General.TypeUnifier.WithVerifiedTypeHandle(RuntimeConstructedGenericTypeInfo, RuntimeTypeInfo[]) + 0x75
   at StreamJsonRpc.SystemTextJsonFormatter.AsyncEnumerableConverter.CreateConverter(Type, JsonSerializerOptions) + 0x84

This code here:

https://github.com/microsoft/vs-streamjsonrpc/blob/81a08406f84ff4f53a2de20dee2637f4e40292a6/src/StreamJsonRpc/SystemTextJsonFormatter.cs#L840-L847

cc @AArnott

[AArnott]

Oh, I see what's going on here. I have that working automatically for the Nerdbank.MessagePack formatter due to some magic we play with the PolyType source generator. This might be a bit more hacky for STJ but I suspect we can come up with something.