Fix race conditions when transitioning instance states during Agent shutdown

Agent/Phantom.Agent.Rpc/RpcServerConnection.cs

@@ -15,29 +15,24 @@ public sealed class RpcServerConnection {
 		this.replyTracker = replyTracker;
 	}
 
-	private byte[] WriteBytes<TMessage, TReply>(TMessage message) where TMessage : IMessageToServer<TReply> {
-		return MessageRegistries.ToServer.Write<TMessage, TReply>(message).ToArray();
-	}
-
 	internal async Task Send<TMessage>(TMessage message) where TMessage : IMessageToServer {
-		var bytes = WriteBytes<TMessage, NoReply>(message);
+		var bytes = MessageRegistries.ToServer.Write(message).ToArray();
 		if (bytes.Length > 0) {
 			await socket.SendAsync(bytes);
 		}
 	}
 
-	internal async Task<TReply?> Send<TMessage, TReply>(Func<uint, TMessage> messageFactory, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToServer<TReply> where TReply : class {
+	internal async Task<TReply?> Send<TMessage, TReply>(TMessage message, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToServer<TReply> where TReply : class {
 		var sequenceId = replyTracker.RegisterReply();
-		var message = messageFactory(sequenceId);
 		
-		var bytes = WriteBytes<TMessage, TReply>(message);
+		var bytes = MessageRegistries.ToServer.Write<TMessage, TReply>(sequenceId, message).ToArray();
 		if (bytes.Length == 0) {
 			replyTracker.ForgetReply(sequenceId);
 			return null;
 		}
 
 		await socket.SendAsync(bytes);
-		return await replyTracker.WaitForReply<TReply>(message.SequenceId, waitForReplyTime, cancellationToken);
+		return await replyTracker.WaitForReply<TReply>(sequenceId, waitForReplyTime, cancellationToken);
 	}
 
 	public void Receive(ReplyMessage message) {

Agent/Phantom.Agent.Rpc/ServerMessaging.cs

@@ -28,7 +28,7 @@ public static class ServerMessaging {
 		return CurrentConnectionOrThrow.Send(message);
 	}
 
-	public static Task<TReply?> Send<TMessage, TReply>(Func<uint, TMessage> messageFactory, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToServer<TReply> where TReply : class {
-		return CurrentConnectionOrThrow.Send<TMessage, TReply>(messageFactory, waitForReplyTime, cancellationToken);
+	public static Task<TReply?> Send<TMessage, TReply>(TMessage message, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToServer<TReply> where TReply : class {
+		return CurrentConnectionOrThrow.Send<TMessage, TReply>(message, waitForReplyTime, cancellationToken);
 	}
 }

Agent/Phantom.Agent.Services/Instances/Instance.cs

@@ -85,10 +85,10 @@ sealed class Instance : IDisposable {
 		}
 	}
 
-	public async Task<LaunchInstanceResult> Launch(CancellationToken cancellationToken) {
-		await stateTransitioningActionSemaphore.WaitAsync(cancellationToken);
+	public async Task<LaunchInstanceResult> Launch(CancellationToken shutdownCancellationToken) {
+		await stateTransitioningActionSemaphore.WaitAsync(shutdownCancellationToken);
 		try {
-			return TransitionStateAndReturn(currentState.Launch(new InstanceContextImpl(this)));
+			return TransitionStateAndReturn(currentState.Launch(new InstanceContextImpl(this, shutdownCancellationToken)));
 		} catch (Exception e) {
 			logger.Error(e, "Caught exception while launching instance.");
 			return LaunchInstanceResult.UnknownError;
@@ -126,10 +126,13 @@ sealed class Instance : IDisposable {
 
 	private sealed class InstanceContextImpl : InstanceContext {
 		private readonly Instance instance;
+		private readonly CancellationToken shutdownCancellationToken;
+		
 		private int statusUpdateCounter;
 
-		public InstanceContextImpl(Instance instance) : base(instance.Configuration, instance.Launcher) {
+		public InstanceContextImpl(Instance instance, CancellationToken shutdownCancellationToken) : base(instance.Configuration, instance.Launcher) {
 			this.instance = instance;
+			this.shutdownCancellationToken = shutdownCancellationToken;
 		}
 
 		public override LaunchServices LaunchServices => instance.launchServices;
@@ -148,10 +151,20 @@ sealed class Instance : IDisposable {
 			});
 		}
 
-		public override void TransitionState(Func<IInstanceState> newState) {
-			instance.stateTransitioningActionSemaphore.Wait();
+		public override void TransitionState(Func<(IInstanceState, IInstanceStatus?)> newStateAndStatus) {
+			instance.stateTransitioningActionSemaphore.Wait(CancellationToken.None);
 			try {
-				instance.TransitionState(newState());
+				var (state, status) = newStateAndStatus();
+				if (state is not InstanceNotRunningState && shutdownCancellationToken.IsCancellationRequested) {
+					instance.logger.Verbose("Cancelled state transition to {State} due to Agent shutdown.", state.GetType().Name);
+					return;
+				}
+
+				if (status != null) {
+					ReportStatus(status);
+				}
+
+				instance.TransitionState(state);
 			} catch (Exception e) {
 				instance.logger.Error(e, "Caught exception during state transition.");
 			} finally {

Agent/Phantom.Agent.Services/Instances/InstanceContext.cs

@@ -20,9 +20,9 @@ abstract class InstanceContext {
 	}
 
 	public abstract void ReportStatus(IInstanceStatus newStatus);
-	public abstract void TransitionState(Func<IInstanceState> newState);
+	public abstract void TransitionState(Func<(IInstanceState, IInstanceStatus?)> newStateAndStatus);
 
-	public void TransitionState(IInstanceState newState) {
-		TransitionState(() => newState);
+	public void TransitionState(IInstanceState newState, IInstanceStatus? newStatus = null) {
+		TransitionState(() => (newState, newStatus));
 	}
 }

Agent/Phantom.Agent.Services/Instances/InstanceSessionManager.cs

@@ -41,36 +41,36 @@ sealed class InstanceSessionManager : IDisposable {
 		this.shutdownCancellationToken = shutdownCancellationTokenSource.Token;
 	}
 
-	[SuppressMessage("ReSharper", "ConvertIfStatementToReturnStatement")]
-	private async Task<InstanceActionResult<T>> AcquireSemaphoreAndRunWithInstance<T>(Guid instanceGuid, Func<Instance, Task<T>> func) {
+	private async Task<InstanceActionResult<T>> AcquireSemaphoreAndRun<T>(Func<Task<InstanceActionResult<T>>> func) {
 		try {
 			await semaphore.WaitAsync(shutdownCancellationToken);
+
+			try {
+				return await func();
+			} finally {
+				semaphore.Release();
+			}
 		} catch (OperationCanceledException) {
 			return InstanceActionResult.General<T>(InstanceActionGeneralResult.AgentShuttingDown);
 		}
+	}
 
-		try {
-			if (!instances.TryGetValue(instanceGuid, out var instance)) {
-				return InstanceActionResult.General<T>(InstanceActionGeneralResult.InstanceDoesNotExist);
-			}
-			else {
+	[SuppressMessage("ReSharper", "ConvertIfStatementToReturnStatement")]
+	private Task<InstanceActionResult<T>> AcquireSemaphoreAndRunWithInstance<T>(Guid instanceGuid, Func<Instance, Task<T>> func) {
+		return AcquireSemaphoreAndRun(async () => {
+			if (instances.TryGetValue(instanceGuid, out var instance)) {
 				return InstanceActionResult.Concrete(await func(instance));
 			}
-		} finally {
-			semaphore.Release();
-		}
+			else {
+				return InstanceActionResult.General<T>(InstanceActionGeneralResult.InstanceDoesNotExist);
+			}
+		});
 	}
 
 	public async Task<InstanceActionResult<ConfigureInstanceResult>> Configure(InstanceConfiguration configuration) {
-		try {
-			await semaphore.WaitAsync(shutdownCancellationToken);
-		} catch (OperationCanceledException) {
-			return InstanceActionResult.General<ConfigureInstanceResult>(InstanceActionGeneralResult.AgentShuttingDown);
-		}
+		return await AcquireSemaphoreAndRun(async () => {
+			var instanceGuid = configuration.InstanceGuid;
 			
-		var instanceGuid = configuration.InstanceGuid;
-
-		try {
 			var otherInstances = instances.Values.Where(inst => inst.Configuration.InstanceGuid != instanceGuid).ToArray();
 			if (otherInstances.Length + 1 > agentInfo.MaxInstances) {
 				return InstanceActionResult.Concrete(ConfigureInstanceResult.InstanceLimitExceeded);
@@ -115,9 +115,7 @@ sealed class InstanceSessionManager : IDisposable {
 			}
 
 			return InstanceActionResult.Concrete(ConfigureInstanceResult.Success);
-		} finally {
-			semaphore.Release();
-		}
+		});
 	}
 
 	public Task<InstanceActionResult<LaunchInstanceResult>> Launch(Guid instanceGuid) {

Agent/Phantom.Agent.Services/Instances/States/InstanceLaunchingState.cs

@@ -66,11 +66,10 @@ sealed class InstanceLaunchingState : IInstanceState, IDisposable {
 		context.TransitionState(() => {
 			if (cancellationTokenSource.IsCancellationRequested) {
 				context.PortManager.Release(context.Configuration);
-				context.ReportStatus(InstanceStatus.NotRunning);
-				return new InstanceNotRunningState();
+				return (new InstanceNotRunningState(), InstanceStatus.NotRunning);
 			}
 			else {
-				return new InstanceRunningState(context, task.Result);
+				return (new InstanceRunningState(context, task.Result), null);
 			}
 		});
 	}

Agent/Phantom.Agent.Services/Instances/States/InstanceRunningState.cs

@@ -30,8 +30,7 @@ sealed class InstanceRunningState : IInstanceState {
 		if (session.HasEnded) {
 			if (sessionObjects.Dispose()) {
 				context.Logger.Warning("Session ended immediately after it was started.");
-				context.ReportStatus(InstanceStatus.Failed(InstanceLaunchFailReason.UnknownError));
-				context.LaunchServices.TaskManager.Run(() => context.TransitionState(new InstanceNotRunningState()));
+				context.LaunchServices.TaskManager.Run(() => context.TransitionState(new InstanceNotRunningState(), InstanceStatus.Failed(InstanceLaunchFailReason.UnknownError)));
 			}
 		}
 		else {
@@ -52,13 +51,11 @@ sealed class InstanceRunningState : IInstanceState {
 
 		if (isStopping) {
 			context.Logger.Information("Session ended.");
-			context.ReportStatus(InstanceStatus.NotRunning);
-			context.TransitionState(new InstanceNotRunningState());
+			context.TransitionState(new InstanceNotRunningState(), InstanceStatus.NotRunning);
 		}
 		else {
 			context.Logger.Information("Session ended unexpectedly, restarting...");
-			context.ReportStatus(InstanceStatus.Restarting);
-			context.TransitionState(new InstanceLaunchingState(context));
+			context.TransitionState(new InstanceLaunchingState(context), InstanceStatus.Restarting);
 		}
 	}
 

Agent/Phantom.Agent.Services/Instances/States/InstanceStoppingState.cs

@@ -1,4 +1,5 @@
-using Phantom.Agent.Minecraft.Command;
+using System.Diagnostics;
+using Phantom.Agent.Minecraft.Command;
 using Phantom.Agent.Minecraft.Instance;
 using Phantom.Common.Data.Instance;
 using Phantom.Common.Data.Minecraft;
@@ -32,26 +33,27 @@ sealed class InstanceStoppingState : IInstanceState, IDisposable {
 			await DoWaitForSessionToEnd();
 		} finally {
 			context.Logger.Information("Session stopped.");
-			context.ReportStatus(InstanceStatus.NotRunning);
-			context.TransitionState(new InstanceNotRunningState());
+			context.TransitionState(new InstanceNotRunningState(), InstanceStatus.NotRunning);
 		}
 	}
 
 	private async Task DoSendStopCommand() {
-		using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+		using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
 		try {
-			await session.SendCommand(MinecraftCommand.Stop, cts.Token);
+			await session.SendCommand(MinecraftCommand.Stop, timeout.Token);
 		} catch (OperationCanceledException) {
 			// ignore
+		} catch (ObjectDisposedException e) when (e.ObjectName == typeof(Process).FullName && session.HasEnded) {
+			// ignore
 		} catch (Exception e) {
 			context.Logger.Warning(e, "Caught exception while sending stop command.");
 		}
 	}
 
 	private async Task DoWaitForSessionToEnd() {
-		using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(55));
+		using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(55));
 		try {
-			await session.WaitForExit(cts.Token);
+			await session.WaitForExit(timeout.Token);
 		} catch (OperationCanceledException) {
 			try {
 				context.Logger.Warning("Waiting timed out, killing session...");

Common/Phantom.Common.Messages/IMessageToAgent.cs

@@ -4,6 +4,4 @@ namespace Phantom.Common.Messages;
 
 public interface IMessageToAgent<TReply> : IMessage<IMessageToAgentListener, TReply> {}
 
-public interface IMessageToAgent : IMessageToAgent<NoReply> {
-	uint IMessage<IMessageToAgentListener, NoReply>.SequenceId => 0;
-}
+public interface IMessageToAgent : IMessageToAgent<NoReply> {}

Common/Phantom.Common.Messages/IMessageToServer.cs

@@ -4,6 +4,4 @@ namespace Phantom.Common.Messages;
 
 public interface IMessageToServer<TReply> : IMessage<IMessageToServerListener, TReply> {}
 
-public interface IMessageToServer : IMessageToServer<NoReply> {
-	uint IMessage<IMessageToServerListener, NoReply>.SequenceId => 0;
-}
+public interface IMessageToServer : IMessageToServer<NoReply> {}

Common/Phantom.Common.Messages/MessageRegistries.cs

@@ -12,20 +12,20 @@ public static class MessageRegistries {
 	public static MessageRegistry<IMessageToServerListener> ToServer { get; } = new (PhantomLogger.Create("MessageRegistry:ToServer"));
 
 	static MessageRegistries() {
-		ToAgent.Add<RegisterAgentSuccessMessage, NoReply>(0);
-		ToAgent.Add<RegisterAgentFailureMessage, NoReply>(1);
+		ToAgent.Add<RegisterAgentSuccessMessage>(0);
+		ToAgent.Add<RegisterAgentFailureMessage>(1);
 		ToAgent.Add<ConfigureInstanceMessage, InstanceActionResult<ConfigureInstanceResult>>(2);
 		ToAgent.Add<LaunchInstanceMessage, InstanceActionResult<LaunchInstanceResult>>(3);
 		ToAgent.Add<StopInstanceMessage, InstanceActionResult<StopInstanceResult>>(4);
 		ToAgent.Add<SendCommandToInstanceMessage, InstanceActionResult<SendCommandToInstanceResult>>(5);
-		ToAgent.Add<ReplyMessage, NoReply>(127);
+		ToAgent.Add<ReplyMessage>(127);
 		
-		ToServer.Add<RegisterAgentMessage, NoReply>(0);
-		ToServer.Add<UnregisterAgentMessage, NoReply>(1);
-		ToServer.Add<AgentIsAliveMessage, NoReply>(2);
-		ToServer.Add<AdvertiseJavaRuntimesMessage, NoReply>(3);
-		ToServer.Add<ReportInstanceStatusMessage, NoReply>(4);
-		ToServer.Add<InstanceOutputMessage, NoReply>(5);
-		ToServer.Add<ReplyMessage, NoReply>(127);
+		ToServer.Add<RegisterAgentMessage>(0);
+		ToServer.Add<UnregisterAgentMessage>(1);
+		ToServer.Add<AgentIsAliveMessage>(2);
+		ToServer.Add<AdvertiseJavaRuntimesMessage>(3);
+		ToServer.Add<ReportInstanceStatusMessage>(4);
+		ToServer.Add<InstanceOutputMessage>(5);
+		ToServer.Add<ReplyMessage>(127);
 	}
 }

Common/Phantom.Common.Messages/ToAgent/ConfigureInstanceMessage.cs

@@ -6,8 +6,7 @@ namespace Phantom.Common.Messages.ToAgent;
 
 [MemoryPackable]
 public sealed partial record ConfigureInstanceMessage(
-	[property: MemoryPackOrder(0)] uint SequenceId,
-	[property: MemoryPackOrder(1)] InstanceConfiguration Configuration
+	[property: MemoryPackOrder(0)] InstanceConfiguration Configuration
 ) : IMessageToAgent<InstanceActionResult<ConfigureInstanceResult>> {
 	public Task<InstanceActionResult<ConfigureInstanceResult>> Accept(IMessageToAgentListener listener) {
 		return listener.HandleConfigureInstance(this);

Common/Phantom.Common.Messages/ToAgent/LaunchInstanceMessage.cs

@@ -5,8 +5,7 @@ namespace Phantom.Common.Messages.ToAgent;
 
 [MemoryPackable]
 public sealed partial record LaunchInstanceMessage(
-	[property: MemoryPackOrder(0)] uint SequenceId,
-	[property: MemoryPackOrder(1)] Guid InstanceGuid
+	[property: MemoryPackOrder(0)] Guid InstanceGuid
 ) : IMessageToAgent<InstanceActionResult<LaunchInstanceResult>> {
 	public Task<InstanceActionResult<LaunchInstanceResult>> Accept(IMessageToAgentListener listener) {
 		return listener.HandleLaunchInstance(this);

Common/Phantom.Common.Messages/ToAgent/SendCommandToInstanceMessage.cs

@@ -5,9 +5,8 @@ namespace Phantom.Common.Messages.ToAgent;
 
 [MemoryPackable]
 public sealed partial record SendCommandToInstanceMessage(
-	[property: MemoryPackOrder(0)] uint SequenceId,
-	[property: MemoryPackOrder(1)] Guid InstanceGuid,
-	[property: MemoryPackOrder(2)] string Command
+	[property: MemoryPackOrder(0)] Guid InstanceGuid,
+	[property: MemoryPackOrder(1)] string Command
 ) : IMessageToAgent<InstanceActionResult<SendCommandToInstanceResult>> {
 	public Task<InstanceActionResult<SendCommandToInstanceResult>> Accept(IMessageToAgentListener listener) {
 		return listener.HandleSendCommandToInstance(this);

Common/Phantom.Common.Messages/ToAgent/StopInstanceMessage.cs

@@ -6,9 +6,8 @@ namespace Phantom.Common.Messages.ToAgent;
 
 [MemoryPackable]
 public sealed partial record StopInstanceMessage(
-	[property: MemoryPackOrder(0)] uint SequenceId,
-	[property: MemoryPackOrder(1)] Guid InstanceGuid,
-	[property: MemoryPackOrder(2)] MinecraftStopStrategy StopStrategy
+	[property: MemoryPackOrder(0)] Guid InstanceGuid,
+	[property: MemoryPackOrder(1)] MinecraftStopStrategy StopStrategy
 ) : IMessageToAgent<InstanceActionResult<StopInstanceResult>> {
 	public Task<InstanceActionResult<StopInstanceResult>> Accept(IMessageToAgentListener listener) {
 		return listener.HandleStopInstance(this);

Server/Phantom.Server.Rpc/RpcClientConnection.cs

@@ -34,29 +34,32 @@ public sealed class RpcClientConnection {
 		}
 	}
 
-	private byte[] WriteBytes<TMessage, TReply>(TMessage message) where TMessage : IMessageToAgent<TReply> {
-		return isClosed ? Array.Empty<byte>() : MessageRegistries.ToAgent.Write<TMessage, TReply>(message).ToArray();
-	}
-
 	public async Task Send<TMessage>(TMessage message) where TMessage : IMessageToAgent {
-		var bytes = WriteBytes<TMessage, NoReply>(message);
+		if (isClosed) {
+			return;
+		}
+		
+		var bytes = MessageRegistries.ToAgent.Write(message).ToArray();
 		if (bytes.Length > 0) {
 			await socket.SendAsync(routingId, bytes);
 		}
 	}
 
-	public async Task<TReply?> Send<TMessage, TReply>(Func<uint, TMessage> messageFactory, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToAgent<TReply> where TReply : class {
-		var sequenceId = messageReplyTracker.RegisterReply();
-		var message = messageFactory(sequenceId);
+	public async Task<TReply?> Send<TMessage, TReply>(TMessage message, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToAgent<TReply> where TReply : class {
+		if (isClosed) {
+			return null;
+		}
 		
-		var bytes = WriteBytes<TMessage, TReply>(message);
+		var sequenceId = messageReplyTracker.RegisterReply();
+		
+		var bytes = MessageRegistries.ToAgent.Write<TMessage, TReply>(sequenceId, message).ToArray();
 		if (bytes.Length == 0) {
 			messageReplyTracker.ForgetReply(sequenceId);
 			return null;
 		}
 
 		await socket.SendAsync(routingId, bytes);
-		return await messageReplyTracker.WaitForReply<TReply>(message.SequenceId, waitForReplyTime, cancellationToken);
+		return await messageReplyTracker.WaitForReply<TReply>(sequenceId, waitForReplyTime, cancellationToken);
 	}
 
 	public void Receive(ReplyMessage message) {

Server/Phantom.Server.Services/Agents/AgentConnection.cs

@@ -22,7 +22,7 @@ sealed class AgentConnection {
 		return connection.Send(message);
 	}
 
-	public Task<TReply?> Send<TMessage, TReply>(Func<uint, TMessage> messageFactory, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToAgent<TReply> where TReply : class {
-		return connection.Send<TMessage, TReply>(messageFactory, waitForReplyTime, cancellationToken);
+	public Task<TReply?> Send<TMessage, TReply>(TMessage message, TimeSpan waitForReplyTime, CancellationToken cancellationToken) where TMessage : IMessageToAgent<TReply> where TReply : class {
+		return connection.Send<TMessage, TReply>(message, waitForReplyTime, cancellationToken);
 	}
 }

Server/Phantom.Server.Services/Agents/AgentManager.cs

@@ -113,14 +113,14 @@ public sealed class AgentManager {
 		}
 	}
 
-	internal async Task<TReply?> SendMessage<TMessage, TReply>(Guid guid, Func<uint, TMessage> messageFactory, TimeSpan waitForReplyTime) where TMessage : IMessageToAgent<TReply> where TReply : class {
+	internal async Task<TReply?> SendMessage<TMessage, TReply>(Guid guid, TMessage message, TimeSpan waitForReplyTime) where TMessage : IMessageToAgent<TReply> where TReply : class {
 		var connection = agents.ByGuid.TryGetValue(guid, out var agent) ? agent.Connection : null;
 		if (connection == null) {
 			// TODO handle missing agent?
 			return null;
 		}
 
-		return await connection.Send<TMessage, TReply>(messageFactory, waitForReplyTime, cancellationToken);
+		return await connection.Send<TMessage, TReply>(message, waitForReplyTime, cancellationToken);
 	}
 
 	private sealed class ObservableAgents : ObservableState<ImmutableArray<Agent>> {

Server/Phantom.Server.Services/Instances/InstanceManager.cs

@@ -3,6 +3,7 @@ using Phantom.Common.Data.Instance;
 using Phantom.Common.Data.Minecraft;
 using Phantom.Common.Data.Replies;
 using Phantom.Common.Logging;
+using Phantom.Common.Messages;
 using Phantom.Common.Messages.ToAgent;
 using Phantom.Common.Minecraft;
 using Phantom.Server.Database;
@@ -67,7 +68,7 @@ public sealed class InstanceManager {
 
 		var agentName = agent.Name;
 
-		var reply = (await agentManager.SendMessage<ConfigureInstanceMessage, InstanceActionResult<ConfigureInstanceResult>>(configuration.AgentGuid, sequenceId => new ConfigureInstanceMessage(sequenceId, configuration), TimeSpan.FromSeconds(10))).DidNotReplyIfNull();
+		var reply = (await agentManager.SendMessage<ConfigureInstanceMessage, InstanceActionResult<ConfigureInstanceResult>>(configuration.AgentGuid, new ConfigureInstanceMessage(configuration), TimeSpan.FromSeconds(10))).DidNotReplyIfNull();
 		if (reply.Is(ConfigureInstanceResult.Success)) {
 			using (var scope = databaseProvider.CreateScope()) {
 				InstanceEntity entity = scope.Ctx.InstanceUpsert.Fetch(configuration.InstanceGuid);
@@ -119,6 +120,11 @@ public sealed class InstanceManager {
 		instances.ByGuid.ReplaceAllIf(instance => instance with { Status = instanceStatus }, instance => instance.Configuration.AgentGuid == agentGuid);
 	}
 
+	private async Task<InstanceActionResult<TReply>> SendInstanceActionMessage<TMessage, TReply>(Instance instance, TMessage message) where TMessage : IMessageToAgent<InstanceActionResult<TReply>> {
+		var reply = await agentManager.SendMessage<TMessage, InstanceActionResult<TReply>>(instance.Configuration.AgentGuid, message, TimeSpan.FromSeconds(10));
+		return reply.DidNotReplyIfNull();
+	}
+
 	public async Task<InstanceActionResult<LaunchInstanceResult>> LaunchInstance(Guid instanceGuid) {
 		var instance = GetInstance(instanceGuid);
 		if (instance == null) {
@@ -127,8 +133,7 @@ public sealed class InstanceManager {
 
 		await SetInstanceShouldLaunchAutomatically(instanceGuid, true);
 
-		var reply = await agentManager.SendMessage<LaunchInstanceMessage, InstanceActionResult<LaunchInstanceResult>>(instance.Configuration.AgentGuid, sequenceId => new LaunchInstanceMessage(sequenceId, instanceGuid), TimeSpan.FromSeconds(10));
-		return reply.DidNotReplyIfNull();
+		return await SendInstanceActionMessage<LaunchInstanceMessage, LaunchInstanceResult>(instance, new LaunchInstanceMessage(instanceGuid));
 	}
 
 	public async Task<InstanceActionResult<StopInstanceResult>> StopInstance(Guid instanceGuid, MinecraftStopStrategy stopStrategy) {
@@ -139,8 +144,7 @@ public sealed class InstanceManager {
 
 		await SetInstanceShouldLaunchAutomatically(instanceGuid, false);
 
-		var reply = await agentManager.SendMessage<StopInstanceMessage, InstanceActionResult<StopInstanceResult>>(instance.Configuration.AgentGuid, sequenceId => new StopInstanceMessage(sequenceId, instanceGuid, stopStrategy), TimeSpan.FromSeconds(10));
-		return reply.DidNotReplyIfNull();
+		return await SendInstanceActionMessage<StopInstanceMessage, StopInstanceResult>(instance, new StopInstanceMessage(instanceGuid, stopStrategy));
 	}
 
 	private async Task SetInstanceShouldLaunchAutomatically(Guid instanceGuid, bool shouldLaunchAutomatically) {
@@ -162,8 +166,7 @@ public sealed class InstanceManager {
 			return InstanceActionResult.General<SendCommandToInstanceResult>(InstanceActionGeneralResult.InstanceDoesNotExist);
 		}
 
-		var reply = await agentManager.SendMessage<SendCommandToInstanceMessage, InstanceActionResult<SendCommandToInstanceResult>>(instance.Configuration.AgentGuid, sequenceId => new SendCommandToInstanceMessage(sequenceId, instanceGuid, command), TimeSpan.FromSeconds(10));
-		return reply.DidNotReplyIfNull();
+		return await SendInstanceActionMessage<SendCommandToInstanceMessage, SendCommandToInstanceResult>(instance, new SendCommandToInstanceMessage(instanceGuid, command));
 	}
 
 	internal ImmutableArray<InstanceConfiguration> GetInstanceConfigurationsForAgent(Guid agentGuid) {

Utils/Phantom.Utils.Rpc/Message/IMessage.cs

@@ -1,6 +1,5 @@
 namespace Phantom.Utils.Rpc.Message; 
 
 public interface IMessage<TListener, TReply> {
-	public uint SequenceId { get; }
 	Task<TReply> Accept(TListener listener);
 }

Utils/Phantom.Utils.Rpc/Message/MessageHandler.cs

@@ -17,22 +17,22 @@ public abstract class MessageHandler<TListener> {
 		this.cancellationToken = cancellationToken;
 	}
 	
-	internal void Enqueue<TMessage, TReply>(TMessage message) where TMessage : IMessage<TListener, TReply> {
+	internal void Enqueue<TMessage, TReply>(uint sequenceId, TMessage message) where TMessage : IMessage<TListener, TReply> {
 		cancellationToken.ThrowIfCancellationRequested();
 		taskManager.Run(async () => {
 			try {
-				await Handle<TMessage, TReply>(message);
+				await Handle<TMessage, TReply>(sequenceId, message);
 			} catch (Exception e) {
 				logger.Error(e, "Failed to handle message {Type}.", message.GetType().Name);
 			}
 		});
 	}
 
-	private async Task Handle<TMessage, TReply>(TMessage message) where TMessage : IMessage<TListener, TReply> {
+	private async Task Handle<TMessage, TReply>(uint sequenceId, TMessage message) where TMessage : IMessage<TListener, TReply> {
 		TReply reply = await message.Accept(Listener);
 		
 		if (reply is not NoReply) {
-			await SendReply(message.SequenceId, MessageSerializer.Serialize(reply));
+			await SendReply(sequenceId, MessageSerializer.Serialize(reply));
 		}
 	}
 	

Utils/Phantom.Utils.Rpc/Message/MessageRegistry.cs

@@ -17,10 +17,23 @@ public sealed class MessageRegistry<TListener> {
 		this.logger = logger;
 	}
 
+	public void Add<TMessage>(ushort code) where TMessage : IMessage<TListener, NoReply> {
+		AddTypeCodeMapping<TMessage, NoReply>(code);
+		codeToHandlerMapping.Add(code, DeserializationHandler<TMessage>);
+	}
+
 	public void Add<TMessage, TReply>(ushort code) where TMessage : IMessage<TListener, TReply> {
+		if (typeof(TReply) == typeof(NoReply)) {
+			throw new InvalidOperationException("This overload of Add must not be used with NoReply as the reply type!");
+		}
+		
+		AddTypeCodeMapping<TMessage, TReply>(code);
+		codeToHandlerMapping.Add(code, DeserializationHandler<TMessage, TReply>);
+	}
+
+	private void AddTypeCodeMapping<TMessage, TReply>(ushort code) where TMessage : IMessage<TListener, TReply> {
 		typeToCodeMapping.Add(typeof(TMessage), code);
 		codeToTypeMapping.Add(code, typeof(TMessage));
-		codeToHandlerMapping.Add(code, HandleInternal<TMessage, TReply>);
 	}
 
 	public bool TryGetType(ReadOnlyMemory<byte> data, [NotNullWhen(true)] out Type? type) {
@@ -34,10 +47,10 @@ public sealed class MessageRegistry<TListener> {
 	}
 
 	public ReadOnlySpan<byte> Write<TMessage>(TMessage message) where TMessage : IMessage<TListener, NoReply> {
-		return Write<TMessage, NoReply>(message);
+		return Write<TMessage, NoReply>(0, message);
 	}
 
-	public ReadOnlySpan<byte> Write<TMessage, TReply>(TMessage message) where TMessage : IMessage<TListener, TReply> {
+	public ReadOnlySpan<byte> Write<TMessage, TReply>(uint sequenceId, TMessage message) where TMessage : IMessage<TListener, TReply> {
 		if (!typeToCodeMapping.TryGetValue(typeof(TMessage), out ushort code)) {
 			logger.Error("Unknown message type {Type}.", typeof(TMessage));
 			return default;
@@ -47,6 +60,11 @@ public sealed class MessageRegistry<TListener> {
 
 		try {
 			MessageSerializer.WriteCode(buffer, code);
+
+			if (typeof(TReply) != typeof(NoReply)) {
+				MessageSerializer.WriteSequenceId(buffer, sequenceId);
+			}
+
 			MessageSerializer.Serialize(buffer, message);
 
 			if (buffer.WrittenCount > DefaultBufferSize && logger.IsEnabled(LogEventLevel.Verbose)) {
@@ -77,7 +95,23 @@ public sealed class MessageRegistry<TListener> {
 		handle(data, code, handler);
 	}
 
-	private void HandleInternal<TMessage, TReply>(ReadOnlyMemory<byte> data, ushort code, MessageHandler<TListener> handler) where TMessage : IMessage<TListener, TReply> {
+	private void DeserializationHandler<TMessage>(ReadOnlyMemory<byte> data, ushort code, MessageHandler<TListener> handler) where TMessage : IMessage<TListener, NoReply> {
+		DeserializeAndEnqueueMessage<TMessage, NoReply>(data, code, handler, 0);
+	}
+
+	private void DeserializationHandler<TMessage, TReply>(ReadOnlyMemory<byte> data, ushort code, MessageHandler<TListener> handler) where TMessage : IMessage<TListener, TReply> {
+		uint sequenceId;
+		try {
+			sequenceId = MessageSerializer.ReadSequenceId(ref data);
+		} catch (Exception e) {
+			logger.Error(e, "Failed to deserialize sequence ID of message with code {Code}.", code);
+			return;
+		}
+
+		DeserializeAndEnqueueMessage<TMessage, TReply>(data, code, handler, sequenceId);
+	}
+
+	private void DeserializeAndEnqueueMessage<TMessage, TReply>(ReadOnlyMemory<byte> data, ushort code, MessageHandler<TListener> handler, uint sequenceId) where TMessage : IMessage<TListener, TReply> {
 		TMessage message;
 		try {
 			message = MessageSerializer.Deserialize<TMessage>(data);
@@ -86,6 +120,6 @@ public sealed class MessageRegistry<TListener> {
 			return;
 		}
 
-		handler.Enqueue<TMessage, TReply>(message);
+		handler.Enqueue<TMessage, TReply>(sequenceId, message);
 	}
 }

Utils/Phantom.Utils.Rpc/Message/MessageSerializer.cs

@@ -30,4 +30,16 @@ static class MessageSerializer {
 		memory = memory[2..];
 		return value;
 	}
+
+	public static void WriteSequenceId(IBufferWriter<byte> destination, uint sequenceId) {
+		Span<byte> buffer = stackalloc byte[4];
+		BinaryPrimitives.WriteUInt32LittleEndian(buffer, sequenceId);
+		destination.Write(buffer);
+	}
+	
+	public static uint ReadSequenceId(ref ReadOnlyMemory<byte> memory) {
+		uint value = BinaryPrimitives.ReadUInt32LittleEndian(memory.Span);
+		memory = memory[4..];
+		return value;
+	}
 }