diff --git a/src/Tgstation.Server.Host/Swarm/SwarmService.cs b/src/Tgstation.Server.Host/Swarm/SwarmService.cs index beddc9e9cb..282a95df32 100644 --- a/src/Tgstation.Server.Host/Swarm/SwarmService.cs +++ b/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 /// /// Helps keep servers connected to the same database in sync by coordinating updates. /// - sealed class SwarmService : ISwarmService, ISwarmOperations, IRestartHandler, IDisposable + sealed class SwarmService : ISwarmService, ISwarmOperations, IDisposable { /// /// Interval at which the swarm controller makes health checks on nodes. @@ -101,11 +102,6 @@ namespace Tgstation.Server.Host.Swarm /// readonly IServerUpdater serverUpdater; - /// - /// The for the . - /// - readonly IRestartRegistration restartRegistration; - /// /// The for the . /// @@ -176,11 +172,6 @@ namespace Tgstation.Server.Host.Swarm /// DateTimeOffset? lastControllerHealthCheck; - /// - /// If was called. - /// - bool restarting; - /// /// If the list has been updated and needs to be resent to clients. /// @@ -193,7 +184,6 @@ namespace Tgstation.Server.Host.Swarm /// The value of . /// The value of . /// The value of . - /// The to register ourselves as a with. /// The value of . /// The value of . /// The containing the value of . @@ -203,7 +193,6 @@ namespace Tgstation.Server.Host.Swarm IDatabaseSeeder databaseSeeder, IAssemblyInformationProvider assemblyInformationProvider, IAbstractHttpClientFactory httpClientFactory, - IServerControl serverControl, IServerUpdater serverUpdater, IAsyncDelayer asyncDelayer, IOptions 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 } /// - public void Dispose() - { - restartRegistration.Dispose(); - serverHealthCheckCancellationTokenSource?.Dispose(); - } + public void Dispose() => serverHealthCheckCancellationTokenSource?.Dispose(); /// 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 /// 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; } - /// - public Task HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken) - { - restarting = !gracefulShutdown; - return Task.CompletedTask; - } - /// public async Task 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; diff --git a/tests/Tgstation.Server.Host.Tests/Swarm/SwarmRpcMapper.cs b/tests/Tgstation.Server.Host.Tests/Swarm/SwarmRpcMapper.cs index afdaf505b5..4892d18b82 100644 --- a/tests/Tgstation.Server.Host.Tests/Swarm/SwarmRpcMapper.cs +++ b/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 incomingRegistrationIds = new(); + + readonly List serverErrors = new(); + + List<(SwarmConfiguration, TestableSwarmNode)> configToNodes; public SwarmRpcMapper(Mock clientMock) { clientMock .Setup(x => x.SendAsync(It.IsNotNull(), It.IsAny())) .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 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(); - 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), controllerMethod.ReturnType); - args.Add(cancellationToken); - var invocationTask = (Task)controllerMethod.Invoke(controller, args.ToArray()); - result = await invocationTask; + var args = new List(); + 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), controllerMethod.ReturnType); + args.Add(cancellationToken); + var invocationTask = (Task)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; } } } diff --git a/tests/Tgstation.Server.Host.Tests/Swarm/TestSwarmProtocol.cs b/tests/Tgstation.Server.Host.Tests/Swarm/TestSwarmProtocol.cs index 24bf0661d3..799c676de9 100644 --- a/tests/Tgstation.Server.Host.Tests/Swarm/TestSwarmProtocol.cs +++ b/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 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 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); } } } diff --git a/tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs b/tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs index dd68830482..fe7ed56bee 100644 --- a/tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs +++ b/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 UpdateTask { get; private set; } + public bool Initialized { get; private set; } + public bool Shutdown { get; private set; } + readonly Mock mockHttpClient; - readonly Mock mockServerControl; readonly Mock mockDBContextFactory; readonly Mock mockDatabaseSeeder; readonly ISetup 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>())) .Callback>((func) => func(mockDatabaseContext)); - mockServerControl = new Mock(); - mockServerControl.Setup(x => x.TryStartUpdate(this, It.IsNotNull())).Returns(TryStartUpdate); - mockServerControl.Setup(x => x.RegisterForRestart(It.IsNotNull())).Returns(Mock.Of()); - var mockHttpClientFactory = new Mock(); mockHttpClient = new Mock(); 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(), It.IsAny())) .Returns( - (delay, ct) => Task.Delay(TimeSpan.FromSeconds(1), ct)); + (delay, ct) => Task.Delay(TimeSpan.FromMilliseconds(100), ct)); var mockServerUpdater = new Mock(); 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()); - Controller = new SwarmController( - Service, - RpcMapper, - mockAssemblyInformationProvider.Object, - mockOptions.Object, - loggerFactory.CreateLogger()); + mockServerUpdater + .Setup(x => x.BeginUpdate(It.IsNotNull(), It.IsNotNull(), It.IsAny())) + .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(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()); + } + + 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 TryInit(bool cancel = false) { @@ -156,14 +221,28 @@ namespace Tgstation.Server.Host.Tests.Swarm return result; } - bool TryStartUpdate(IServerUpdateExecutor updateExecutor, Version newVersion) + Task 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 ExecuteUpdate(string updatePath, CancellationToken cancellationToken, CancellationToken criticalCancellationToken) + async Task 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; } } }