Tested swarm controller and fixed some edge cases

- No longer registers for restart unnecessarily
- Fixed infinite wait when committing update
- Fixed unregistering without a registration
- Abort updates if the controller health check fails

src/Tgstation.Server.Host/Swarm/SwarmService.cs

@@ -18,6 +18,7 @@ using Tgstation.Server.Api.Models.Response;
 using Tgstation.Server.Host.Configuration;
 using Tgstation.Server.Host.Core;
 using Tgstation.Server.Host.Database;
+using Tgstation.Server.Host.Extensions;
 using Tgstation.Server.Host.Extensions.Converters;
 using Tgstation.Server.Host.System;
 
@@ -26,7 +27,7 @@ namespace Tgstation.Server.Host.Swarm
 	/// <summary>
 	/// Helps keep servers connected to the same database in sync by coordinating updates.
 	/// </summary>
-	sealed class SwarmService : ISwarmService, ISwarmOperations, IRestartHandler, IDisposable
+	sealed class SwarmService : ISwarmService, ISwarmOperations, IDisposable
 	{
 		/// <summary>
 		/// Interval at which the swarm controller makes health checks on nodes.
@@ -101,11 +102,6 @@ namespace Tgstation.Server.Host.Swarm
 		/// </summary>
 		readonly IServerUpdater serverUpdater;
 
-		/// <summary>
-		/// The <see cref="IRestartRegistration"/> for the <see cref="SwarmService"/>.
-		/// </summary>
-		readonly IRestartRegistration restartRegistration;
-
 		/// <summary>
 		/// The <see cref="ILogger"/> for the <see cref="SwarmService"/>.
 		/// </summary>
@@ -176,11 +172,6 @@ namespace Tgstation.Server.Host.Swarm
 		/// </summary>
 		DateTimeOffset? lastControllerHealthCheck;
 
-		/// <summary>
-		/// If <see cref="IRestartHandler.HandleRestart(Version, bool, CancellationToken)"/> was called.
-		/// </summary>
-		bool restarting;
-
 		/// <summary>
 		/// If the <see cref="swarmServers"/> list has been updated and needs to be resent to clients.
 		/// </summary>
@@ -193,7 +184,6 @@ namespace Tgstation.Server.Host.Swarm
 		/// <param name="databaseSeeder">The value of <see cref="databaseSeeder"/>.</param>
 		/// <param name="assemblyInformationProvider">The value of <see cref="assemblyInformationProvider"/>.</param>
 		/// <param name="httpClientFactory">The value of <see cref="httpClientFactory"/>.</param>
-		/// <param name="serverControl">The <see cref="IServerControl"/> to register ourselves as a <see cref="IRestartHandler"/> with.</param>
 		/// <param name="serverUpdater">The value of <see cref="serverUpdater"/>.</param>
 		/// <param name="asyncDelayer">The value of <see cref="asyncDelayer"/>.</param>
 		/// <param name="swarmConfigurationOptions">The <see cref="IOptions{TOptions}"/> containing the value of <see cref="swarmConfiguration"/>.</param>
@@ -203,7 +193,6 @@ namespace Tgstation.Server.Host.Swarm
 			IDatabaseSeeder databaseSeeder,
 			IAssemblyInformationProvider assemblyInformationProvider,
 			IAbstractHttpClientFactory httpClientFactory,
-			IServerControl serverControl,
 			IServerUpdater serverUpdater,
 			IAsyncDelayer asyncDelayer,
 			IOptions<SwarmConfiguration> swarmConfigurationOptions,
@@ -213,10 +202,6 @@ namespace Tgstation.Server.Host.Swarm
 			this.databaseSeeder = databaseSeeder ?? throw new ArgumentNullException(nameof(databaseSeeder));
 			this.assemblyInformationProvider = assemblyInformationProvider ?? throw new ArgumentNullException(nameof(assemblyInformationProvider));
 			this.httpClientFactory = httpClientFactory ?? throw new ArgumentNullException(nameof(httpClientFactory));
-			if (serverControl == null)
-				throw new ArgumentNullException(nameof(serverControl));
-			restartRegistration = serverControl.RegisterForRestart(this);
-
 			this.serverUpdater = serverUpdater ?? throw new ArgumentNullException(nameof(serverUpdater));
 			this.asyncDelayer = asyncDelayer ?? throw new ArgumentNullException(nameof(asyncDelayer));
 			swarmConfiguration = swarmConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(swarmConfigurationOptions));
@@ -253,11 +238,7 @@ namespace Tgstation.Server.Host.Swarm
 		}
 
 		/// <inheritdoc />
-		public void Dispose()
-		{
-			restartRegistration.Dispose();
-			serverHealthCheckCancellationTokenSource?.Dispose();
-		}
+		public void Dispose() => serverHealthCheckCancellationTokenSource?.Dispose();
 
 		/// <inheritdoc />
 		public async Task RemoteAbortUpdate(CancellationToken cancellationToken)
@@ -278,8 +259,9 @@ namespace Tgstation.Server.Host.Swarm
 				return;
 
 			logger.LogInformation("Aborting swarm update!");
-			updateCommitTcs?.TrySetResult(false);
+			var commitTcs = updateCommitTcs;
 			updateCommitTcs = null;
+			commitTcs?.TrySetResult(false);
 			nodesThatNeedToBeReadyToCommit = null;
 			targetUpdateVersion = null;
 
@@ -332,6 +314,15 @@ namespace Tgstation.Server.Host.Swarm
 			if (!SwarmMode)
 				return SwarmCommitResult.ContinueUpdateNonCommitted;
 
+			// wait for the update commit TCS
+			var commitTcsTask = updateCommitTcs?.Task;
+			if (commitTcsTask == null)
+			{
+				logger.LogDebug("Update commit failed, no pending task completion source!");
+				await AbortUpdate(cancellationToken);
+				return SwarmCommitResult.AbortUpdate;
+			}
+
 			logger.LogInformation("Waiting to commit update...");
 			using var httpClient = httpClientFactory.CreateClient();
 			if (!swarmController)
@@ -357,20 +348,11 @@ namespace Tgstation.Server.Host.Swarm
 				}
 			}
 
-			// wait for the update commit TCS
-			var commitTcsTask = updateCommitTcs?.Task;
-			if (commitTcsTask == null)
-			{
-				logger.LogDebug("Update commit failed, no pending task completion source!");
-				await AbortUpdate(cancellationToken);
-				return SwarmCommitResult.AbortUpdate;
-			}
-
 			var timeoutTask = swarmController
 				? asyncDelayer.Delay(
 					TimeSpan.FromMinutes(UpdateCommitTimeoutMinutes),
 					cancellationToken)
-				: Extensions.TaskExtensions.InfiniteTask();
+				: Extensions.TaskExtensions.InfiniteTask().WithToken(cancellationToken);
 
 			var commitTask = Task.WhenAny(commitTcsTask, timeoutTask);
 
@@ -491,6 +473,8 @@ namespace Tgstation.Server.Host.Swarm
 		/// <inheritdoc />
 		public async Task Shutdown(CancellationToken cancellationToken)
 		{
+			logger.LogTrace("Begin Shutdown");
+
 			async Task SendUnregistrationRequest(SwarmServerResponse swarmServer)
 			{
 				using var httpClient = httpClientFactory.CreateClient();
@@ -524,15 +508,11 @@ namespace Tgstation.Server.Host.Swarm
 
 			if (!swarmController)
 			{
-				// if we restart a node, we don't want to unregister it so the controller doesn't try to update without it
-				// if we're shutting it down, though we should unregister it
-				if (!restarting)
+				if (controllerRegistration != null)
 				{
 					logger.LogInformation("Unregistering from swarm controller...");
 					await SendUnregistrationRequest(null);
 				}
-				else
-					logger.LogTrace("Not unregistering from swarm controller as we are restarting");
 
 				return;
 			}
@@ -541,8 +521,7 @@ namespace Tgstation.Server.Host.Swarm
 			if (targetUpdateVersion != null
 				&& targetUpdateVersion < assemblyInformationProvider.Version)
 				await databaseContextFactory.UseContext(
-					db => databaseSeeder.Downgrade(db, targetUpdateVersion, cancellationToken))
-					;
+					db => databaseSeeder.Downgrade(db, targetUpdateVersion, cancellationToken));
 
 			if (SwarmMode)
 			{
@@ -663,13 +642,6 @@ namespace Tgstation.Server.Host.Swarm
 			return true;
 		}
 
-		/// <inheritdoc />
-		public Task HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken)
-		{
-			restarting = !gracefulShutdown;
-			return Task.CompletedTask;
-		}
-
 		/// <inheritdoc />
 		public async Task<bool> RemoteCommitRecieved(Guid registrationId, CancellationToken cancellationToken)
 		{
@@ -705,8 +677,7 @@ namespace Tgstation.Server.Host.Swarm
 				{
 					logger.LogTrace("All nodes ready, update commit is a go once controller is ready");
 					var commitTcs = updateCommitTcs;
-					commitTcs?.TrySetResult(true);
-					return commitTcs != null;
+					return commitTcs?.TrySetResult(true) == true;
 				}
 			}
 
@@ -1017,6 +988,7 @@ namespace Tgstation.Server.Host.Swarm
 				{
 					logger.LogWarning(ex, "Error during swarm controller health check! Attempting to re-register...");
 					controllerRegistration = null;
+					await AbortUpdate(cancellationToken);
 				}
 
 			SwarmRegistrationResult registrationResult;

tests/Tgstation.Server.Host.Tests/Swarm/SwarmRpcMapper.cs

@@ -19,40 +19,49 @@ using Newtonsoft.Json;
 using Tgstation.Server.Common;
 using Tgstation.Server.Host.Configuration;
 using Tgstation.Server.Host.Controllers;
+using Tgstation.Server.Host.Core;
 using Tgstation.Server.Host.Swarm;
 
 namespace Tgstation.Server.Host.Tests.Swarm
 {
-	sealed class SwarmRpcMapper : IRequestSwarmRegistrationParser
+	sealed class SwarmRpcMapper : IRequestSwarmRegistrationParser, IDisposable
 	{
-		List<(SwarmConfiguration, TestableSwarmNode)> configToControllers;
-		Guid? incomingRegistrationId;
+		public bool AsyncRequests { get; set; }
+
+		readonly Stack<Guid> incomingRegistrationIds = new();
+
+		readonly List<Exception> serverErrors = new();
+
+		List<(SwarmConfiguration, TestableSwarmNode)> configToNodes;
 
 		public SwarmRpcMapper(Mock<IHttpClient> clientMock)
 		{
 			clientMock
 				.Setup(x => x.SendAsync(It.IsNotNull<HttpRequestMessage>(), It.IsAny<CancellationToken>()))
 				.Returns(MapRequest);
+			AsyncRequests = true;
 		}
 
-		public Guid GetRequestRegistrationId(HttpRequest request)
+		public void Dispose()
 		{
-			Assert.IsTrue(incomingRegistrationId.HasValue);
-			var result = incomingRegistrationId.Value;
-			incomingRegistrationId = null;
-			return result;
+			if (serverErrors.Count > 1)
+				throw new AggregateException(serverErrors);
+			else if (serverErrors.Count == 1)
+				throw serverErrors[0];
 		}
 
-		public void Register(List<(SwarmConfiguration, TestableSwarmNode)> configToControllers)
+		public Guid GetRequestRegistrationId(HttpRequest request) => incomingRegistrationIds.Peek();
+
+		public void Register(List<(SwarmConfiguration, TestableSwarmNode)> configToNodes)
 		{
-			this.configToControllers = configToControllers;
+			this.configToNodes = configToNodes;
 		}
 
 		async Task<HttpResponseMessage> MapRequest(
 			HttpRequestMessage request,
 			CancellationToken cancellationToken)
 		{
-			var (config, node) = configToControllers.FirstOrDefault(
+			var (config, node) = configToNodes.FirstOrDefault(
 				pair => pair.Item1.Address.IsBaseOf(request.RequestUri));
 
 			if (config == default)
@@ -63,6 +72,11 @@ namespace Tgstation.Server.Host.Tests.Swarm
 				throw new HttpRequestException("Can't connect to uninitialized node!");
 			}
 
+			if (node.Shutdown)
+			{
+				throw new HttpRequestException("Can't connect to shutdown node!");
+			}
+
 			var controller = node.Controller;
 
 			Type targetAttribute = null;
@@ -106,7 +120,7 @@ namespace Tgstation.Server.Host.Tests.Swarm
 				.Where(pair => pair.Item2 != null
 					&& pair.Item2.HttpMethods.Count() == 1
 					&& pair.Item2.HttpMethods.All(supportedMethod => supportedMethod.Equals(request.Method.Method))
-					&& pair.Item2.Template == route)
+					&& (pair.Item2.Template ?? String.Empty) == route)
 				.Select(pair => pair.method)
 				.SingleOrDefault();
 
@@ -114,63 +128,76 @@ namespace Tgstation.Server.Host.Tests.Swarm
 				Assert.Fail($"SwarmController has no method with attribute {targetAttribute}!");
 
 			// We're not testing OnActionExecutingAsync, that's covered by integration.
-			if (request.Headers.TryGetValues(SwarmConstants.RegistrationIdHeader, out var values) && values.Count() == 1)
-				node.RpcMapper.incomingRegistrationId = Guid.Parse(values.First());
-
-			var args = new List<object>();
-			if (isDataRequest)
-			{
-				var dataType = controllerMethod.GetParameters().First().ParameterType;
-				var json = await request.Content.ReadAsStringAsync(cancellationToken);
-				var parameter = JsonConvert.DeserializeObject(json, dataType, SwarmService.SerializerSettings);
-				args.Add(parameter);
-			}
+			var hasRegistrationHeader = request.Headers.TryGetValues(SwarmConstants.RegistrationIdHeader, out var values) && values.Count() == 1;
+			if (hasRegistrationHeader)
+				node.RpcMapper.incomingRegistrationIds.Push(Guid.Parse(values.First()));
 
 			IActionResult result;
-
-			var response = new HttpResponseMessage();
 			try
 			{
-				if (controllerMethod.ReturnType != typeof(IActionResult))
+				var response = new HttpResponseMessage();
+				try
 				{
-					Assert.AreEqual(typeof(Task<IActionResult>), controllerMethod.ReturnType);
-					args.Add(cancellationToken);
-					var invocationTask = (Task<IActionResult>)controllerMethod.Invoke(controller, args.ToArray());
-					result = await invocationTask;
+					var args = new List<object>();
+					if (isDataRequest && request.Content != null)
+					{
+						var dataType = controllerMethod.GetParameters().First().ParameterType;
+						var json = await request.Content.ReadAsStringAsync(cancellationToken);
+						var parameter = JsonConvert.DeserializeObject(json, dataType, SwarmService.SerializerSettings);
+						args.Add(parameter);
+					}
+
+					if (AsyncRequests)
+						await Task.Yield();
+
+					if (controllerMethod.ReturnType != typeof(IActionResult))
+					{
+						Assert.AreEqual(typeof(Task<IActionResult>), controllerMethod.ReturnType);
+						args.Add(cancellationToken);
+						var invocationTask = (Task<IActionResult>)controllerMethod.Invoke(controller, args.ToArray());
+						result = await invocationTask;
+					}
+					else
+					{
+						result = (IActionResult)controllerMethod.Invoke(controller, args.ToArray());
+
+						// simulate worst case, request completed but was aborted before server replied
+						cancellationToken.ThrowIfCancellationRequested();
+					}
 				}
+				catch (Exception ex)
+				{
+					if (ex is not OperationCanceledException)
+						serverErrors.Add(ex);
+					response.Dispose();
+					throw;
+				}
+
+				// manually checked all controller response types
+				// Fobid, NoContent, Conflict, StatusCode
+				if (result is ForbidResult forbidResult)
+					response.StatusCode = HttpStatusCode.Forbidden;
+				else if (result is NoContentResult noContentResult)
+					response.StatusCode = (HttpStatusCode)noContentResult.StatusCode;
+				else if (result is ConflictResult conflictResult)
+					response.StatusCode = (HttpStatusCode)conflictResult.StatusCode;
+				else if (result is ObjectResult objectResult)
+					response.StatusCode = (HttpStatusCode)objectResult.StatusCode;
+				else if (result is StatusCodeResult statusCodeResult)
+					response.StatusCode = (HttpStatusCode)statusCodeResult.StatusCode;
 				else
 				{
-					result = (IActionResult)controllerMethod.Invoke(controller, args.ToArray());
-
-					// simulate worst case, request completed but was aborted before server replied
-					cancellationToken.ThrowIfCancellationRequested();
+					response.Dispose();
+					Assert.Fail($"Unrecognized result type: {result.GetType()}");
 				}
-			}
-			catch
-			{
-				response.Dispose();
-				throw;
-			}
 
-			// manually checked all controller response types
-			// Fobid, NoContent, Conflict, StatusCode
-			if (result is ForbidResult forbidResult)
-				response.StatusCode = HttpStatusCode.Forbidden;
-			else if (result is NoContentResult noContentResult)
-				response.StatusCode = (HttpStatusCode)noContentResult.StatusCode;
-			else if (result is ConflictResult conflictResult)
-				response.StatusCode = (HttpStatusCode)conflictResult.StatusCode;
-			else if (result is ObjectResult objectResult)
-				response.StatusCode = (HttpStatusCode)objectResult.StatusCode;
-			else if (result is StatusCodeResult statusCodeResult)
-				response.StatusCode = (HttpStatusCode)statusCodeResult.StatusCode;
-			else
-			{
-				response.Dispose();
-				Assert.Fail($"Unrecognized result type: {result.GetType()}");
+				return response;
+			}
+			finally
+			{
+				if (hasRegistrationHeader)
+					node.RpcMapper.incomingRegistrationIds.Pop();
 			}
-
-			return response;
 		}
 	}
 }

tests/Tgstation.Server.Host.Tests/Swarm/TestSwarmProtocol.cs

@@ -13,7 +13,11 @@ namespace Tgstation.Server.Host.Tests.Swarm
 	[TestClass]
 	public sealed class TestSwarmProtocol : IDisposable
 	{
-		readonly ILoggerFactory loggerFactory = LoggerFactory.Create(builder => builder.AddConsole());
+		readonly ILoggerFactory loggerFactory = LoggerFactory.Create(builder =>
+		{
+			builder.SetMinimumLevel(LogLevel.Trace);
+			builder.AddConsole();
+		});
 		readonly HashSet<ushort> usedPorts = new ();
 
 		public void Dispose() => loggerFactory.Dispose();
@@ -27,11 +31,14 @@ namespace Tgstation.Server.Host.Tests.Swarm
 		[TestMethod]
 		public async Task TestInitHappensInstantlyWhenControllerIsInitialized()
 		{
-			using var controller = new TestableSwarmNode(loggerFactory, GenConfig());
-			using var node = new TestableSwarmNode(loggerFactory, GenConfig(controller));
+			await using var controller = GenNode();
+			await using var node = GenNode(controller);
 
 			TestableSwarmNode.Link(controller, node);
 
+			controller.RpcMapper.AsyncRequests = false;
+			node.RpcMapper.AsyncRequests = false;
+
 			var controllerInit = controller.TryInit();
 			Assert.IsTrue(controllerInit.IsCompleted);
 			Assert.AreEqual(SwarmRegistrationResult.Success, await controllerInit);
@@ -44,9 +51,9 @@ namespace Tgstation.Server.Host.Tests.Swarm
 		[TestMethod]
 		public async Task TestNodeInitializeDoesNotWorkWithoutController()
 		{
-			using var controller = new TestableSwarmNode(loggerFactory, GenConfig());
-			using var node1 = new TestableSwarmNode(loggerFactory, GenConfig(controller));
-			using var node2 = new TestableSwarmNode(loggerFactory, GenConfig(controller));
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+			await using var node2 = GenNode(controller);
 
 			TestableSwarmNode.Link(controller, node1, node2);
 
@@ -57,7 +64,155 @@ namespace Tgstation.Server.Host.Tests.Swarm
 			Assert.AreEqual(SwarmRegistrationResult.Success, await node2.TryInit());
 		}
 
-		SwarmConfiguration GenConfig(TestableSwarmNode controller = null)
+		[TestMethod]
+		public async Task TestUpdateCantProceedWithoutAllNodes()
+		{
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+			await using var node2 = GenNode(controller);
+
+			TestableSwarmNode.Link(controller, node1, node2);
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node2.TryInit());
+
+			Assert.IsFalse(await controller.Service.PrepareUpdate(new Version(4, 3, 2), default));
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Assert.IsTrue(await controller.Service.PrepareUpdate(new Version(4, 3, 2), default));
+		}
+
+		[TestMethod]
+		public async Task TestServersReconnectAfterReboot()
+		{
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+
+			TestableSwarmNode.Link(controller, node1);
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
+			Assert.AreEqual(1, node1.Service.GetSwarmServers().Count);
+
+			await DelayMax(() => Assert.AreEqual(2, node1.Service.GetSwarmServers().Count));
+
+			await controller.SimulateReboot(default);
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(1, controller.Service.GetSwarmServers().Count);
+
+			// Remember node's async delayer fires every millisecond
+			// server list is updated in health check loop
+			await DelayMax(() =>
+			{
+				Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
+				Assert.AreEqual(2, node1.Service.GetSwarmServers().Count);
+			});
+
+			await node1.SimulateReboot(default);
+			// health check timeout
+			await DelayMax(() => Assert.AreEqual(1, controller.Service.GetSwarmServers().Count));
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Assert.AreEqual(2, controller.Service.GetSwarmServers().Count);
+			await DelayMax(() => Assert.AreEqual(2, node1.Service.GetSwarmServers().Count));
+		}
+
+		[TestMethod]
+		public async Task TestCommitFailsIfNotPrepared()
+		{
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+
+			TestableSwarmNode.Link(controller, node1);
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Assert.AreEqual(SwarmCommitResult.AbortUpdate, await node1.Service.CommitUpdate(default));
+		}
+
+		[TestMethod]
+		public async Task TestPrepareDifferentVersionsFails()
+		{
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+
+			TestableSwarmNode.Link(controller, node1);
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Assert.IsTrue(await node1.Service.PrepareUpdate(new Version(4, 3, 2), default));
+			Assert.IsFalse(await node1.Service.PrepareUpdate(new Version(4, 3, 3), default));
+		}
+
+		[TestMethod]
+		public async Task TestSimultaneousPrepareDifferentVersionsFailsControllerFirst()
+		{
+			await TestSimultaneousPrepareDifferentVersionsFails(true);
+		}
+
+		[TestMethod]
+		public async Task TestSimultaneousPrepareDifferentVersionsFailsNodeFirst()
+		{
+			await TestSimultaneousPrepareDifferentVersionsFails(false);
+		}
+
+		async Task TestSimultaneousPrepareDifferentVersionsFails(bool prepControllerFirst)
+		{
+			await using var controller = GenNode();
+			await using var node1 = GenNode(controller);
+
+			TestableSwarmNode.Link(controller, node1);
+
+			Assert.AreEqual(SwarmRegistrationResult.Success, await controller.TryInit());
+			Assert.AreEqual(SwarmRegistrationResult.Success, await node1.TryInit());
+
+			Task<bool> controllerPrepareTask = null, nodePrepareTask;
+			if (prepControllerFirst)
+			{
+				controllerPrepareTask = controller.Service.PrepareUpdate(new Version(4, 3, 2), default);
+				Assert.IsFalse(controllerPrepareTask.IsCompleted);
+			}
+
+			nodePrepareTask = node1.Service.PrepareUpdate(new Version(4, 3, 3), default);
+			Assert.IsFalse(nodePrepareTask.IsCompleted);
+
+			if (!prepControllerFirst)
+			{
+				controllerPrepareTask = controller.Service.PrepareUpdate(new Version(4, 3, 2), default);
+				Assert.IsFalse(controllerPrepareTask.IsCompleted);
+			}
+
+			await Task.Yield();
+
+			var controllerResult = await controllerPrepareTask;
+			var nodeResult = await nodePrepareTask;
+
+				Assert.IsFalse(controllerResult && nodeResult);
+		}
+
+		static async Task DelayMax(Action assertion, ulong seconds = 1)
+		{
+			for (var i = 0U; i < (seconds * 10); ++i)
+			{
+				try
+				{
+					assertion();
+					return;
+				}
+				catch (AssertFailedException) { }
+				await Task.Delay(TimeSpan.FromMilliseconds(100));
+			}
+
+			assertion();
+		}
+
+		TestableSwarmNode GenNode(TestableSwarmNode controller = null, Version version = null)
 		{
 			ushort randPort;
 			do
@@ -66,7 +221,7 @@ namespace Tgstation.Server.Host.Tests.Swarm
 			}
 			while (randPort == 0 || !usedPorts.Add(randPort));
 
-			return new SwarmConfiguration
+			var config = new SwarmConfiguration
 			{
 				Address = new Uri($"http://127.0.0.1:{randPort}"),
 				ControllerAddress = controller?.Config.Address,
@@ -74,6 +229,8 @@ namespace Tgstation.Server.Host.Tests.Swarm
 				PrivateKey = "asdf",
 				UpdateRequiredNodeCount = (uint)usedPorts.Count - 1,
 			};
+
+			return new TestableSwarmNode(loggerFactory, config, version);
 		}
 	}
 }

tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Linq;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -20,24 +21,34 @@ using Tgstation.Server.Host.System;
 
 namespace Tgstation.Server.Host.Tests.Swarm
 {
-	sealed class TestableSwarmNode : IDisposable, IServerUpdateExecutor
+	sealed class TestableSwarmNode : IAsyncDisposable
 	{
-		public SwarmController Controller { get; }
+		public SwarmController Controller { get; private set; }
 
-		public SwarmService Service { get; }
+		public SwarmService Service { get; private set; }
 
 		public SwarmConfiguration Config { get; }
 
 		public SwarmRpcMapper RpcMapper { get; }
 
+		public bool UpdateCommits { get; set; }
+
+		public CancellationTokenSource CriticalCancellationTokenSource { get; private set; }
+
+		public ServerUpdateResult UpdateResult { get; set; }
+		public Task<SwarmCommitResult?> UpdateTask { get; private set; }
+
 		public bool Initialized { get; private set; }
 
+		public bool Shutdown { get; private set; }
+
 		readonly Mock<IHttpClient> mockHttpClient;
-		readonly Mock<IServerControl> mockServerControl;
 		readonly Mock<IDatabaseContextFactory> mockDBContextFactory;
 		readonly Mock<IDatabaseSeeder> mockDatabaseSeeder;
 		readonly ISetup<IDatabaseSeeder, Task> mockDatabaseSeederInitialize;
 
+		readonly Action recreateControllerAndService;
+
 		public static void Link(params TestableSwarmNode[] nodes)
 		{
 			var configControllerSet = nodes.Select(x => (x.Config, x)).ToList();
@@ -82,10 +93,6 @@ namespace Tgstation.Server.Host.Tests.Swarm
 				.Setup(x => x.UseContext(It.IsNotNull<Func<IDatabaseContext, Task>>()))
 				.Callback<Func<IDatabaseContext, Task>>((func) => func(mockDatabaseContext));
 
-			mockServerControl = new Mock<IServerControl>();
-			mockServerControl.Setup(x => x.TryStartUpdate(this, It.IsNotNull<Version>())).Returns(TryStartUpdate);
-			mockServerControl.Setup(x => x.RegisterForRestart(It.IsNotNull<IRestartHandler>())).Returns(Mock.Of<IRestartRegistration>());
-
 			var mockHttpClientFactory = new Mock<IAbstractHttpClientFactory>();
 			mockHttpClient = new Mock<IHttpClient>();
 			mockHttpClientFactory.Setup(x => x.CreateClient()).Returns(mockHttpClient.Object);
@@ -94,32 +101,90 @@ namespace Tgstation.Server.Host.Tests.Swarm
 			mockAsyncDelayer.Setup(
 				x => x.Delay(It.IsAny<TimeSpan>(), It.IsAny<CancellationToken>()))
 				.Returns<TimeSpan, CancellationToken>(
-					(delay, ct) => Task.Delay(TimeSpan.FromSeconds(1), ct));
+					(delay, ct) => Task.Delay(TimeSpan.FromMilliseconds(100), ct));
 
 			var mockServerUpdater = new Mock<IServerUpdater>();
 
 			RpcMapper = new SwarmRpcMapper(mockHttpClient);
 
-			Service = new SwarmService(
-				mockDBContextFactory.Object,
-				mockDatabaseSeeder.Object,
-				mockAssemblyInformationProvider.Object,
-				mockHttpClientFactory.Object,
-				mockServerControl.Object,
-				mockServerUpdater.Object,
-				mockAsyncDelayer.Object,
-				mockOptions.Object,
-				loggerFactory.CreateLogger<SwarmService>());
 
-			Controller = new SwarmController(
-				Service,
-				RpcMapper,
-				mockAssemblyInformationProvider.Object,
-				mockOptions.Object,
-				loggerFactory.CreateLogger<SwarmController>());
+			mockServerUpdater
+				.Setup(x => x.BeginUpdate(It.IsNotNull<SwarmService>(), It.IsNotNull<Version>(), It.IsAny<CancellationToken>()))
+				.Returns(BeginUpdate);
+
+			UpdateResult = ServerUpdateResult.Started;
+			UpdateCommits = true;
+
+
+			var runCount = 0;
+			void RecreateControllerAndService()
+			{
+				var run = ++runCount;
+				Initialized = false;
+				Shutdown = false;
+
+				CriticalCancellationTokenSource?.Dispose();
+				CriticalCancellationTokenSource = new CancellationTokenSource();
+
+				var logger = new Logger<SwarmService>(loggerFactory);
+				// HAX HAX HAX
+				logger
+					.GetType()
+					.GetField("_logger", BindingFlags.NonPublic | BindingFlags.Instance)
+					.SetValue(logger, loggerFactory.CreateLogger($"SwarmService-{swarmConfiguration.Identifier}{(run != 1 ? $"-Run{run}": String.Empty)}"));
+
+				Service = new SwarmService(
+					mockDBContextFactory.Object,
+					mockDatabaseSeeder.Object,
+					mockAssemblyInformationProvider.Object,
+					mockHttpClientFactory.Object,
+					mockServerUpdater.Object,
+					mockAsyncDelayer.Object,
+					mockOptions.Object,
+					logger);
+
+				Controller = new SwarmController(
+					Service,
+					RpcMapper,
+					mockAssemblyInformationProvider.Object,
+					mockOptions.Object,
+					loggerFactory.CreateLogger<SwarmController>());
+			}
+
+			RecreateControllerAndService();
+			recreateControllerAndService = RecreateControllerAndService;
 		}
 
-		public void Dispose() => Service.Dispose();
+		public async Task SimulateReboot(CancellationToken cancellationToken)
+		{
+			await ShutdownService(cancellationToken);
+			recreateControllerAndService();
+		}
+
+		public async ValueTask DisposeAsync()
+		{
+			await ShutdownService(default);
+			RpcMapper.Dispose();
+			CriticalCancellationTokenSource.Dispose();
+		}
+
+		private async Task ShutdownService(CancellationToken cancellationToken)
+		{
+			Shutdown = true;
+			CriticalCancellationTokenSource.Cancel();
+			if (UpdateTask != null)
+				try
+				{
+					await UpdateTask;
+				}
+				catch (OperationCanceledException)
+				{
+				}
+			UpdateTask = null;
+
+			await Service.Shutdown(cancellationToken);
+			Service.Dispose();
+		}
 
 		public async Task<SwarmRegistrationResult?> TryInit(bool cancel = false)
 		{
@@ -156,14 +221,28 @@ namespace Tgstation.Server.Host.Tests.Swarm
 			return result;
 		}
 
-		bool TryStartUpdate(IServerUpdateExecutor updateExecutor, Version newVersion)
+		Task<ServerUpdateResult> BeginUpdate(ISwarmService swarmService, Version version, CancellationToken cancellationToken)
 		{
-			throw new NotImplementedException();
+			if (UpdateResult == ServerUpdateResult.Started)
+			{
+				UpdateTask = ExecuteUpdate(version, cancellationToken, CriticalCancellationTokenSource.Token);
+			}
+
+			return Task.FromResult(UpdateResult);
 		}
 
-		public Task<bool> ExecuteUpdate(string updatePath, CancellationToken cancellationToken, CancellationToken criticalCancellationToken)
+		async Task<SwarmCommitResult?> ExecuteUpdate(Version version, CancellationToken cancellationToken, CancellationToken criticalCancellationToken)
 		{
-			throw new NotImplementedException();
+			await Task.Yield(); // Important to simulate some actual kind of asyncronicity here
+
+			await Service.PrepareUpdate(version, cancellationToken);
+
+			if (UpdateCommits)
+			{
+				return await Service.CommitUpdate(criticalCancellationToken);
+			}
+
+			return null;
 		}
 	}
 }