From 4a6b549a790dbf3e8499ae9df89ffe5db74eea7f Mon Sep 17 00:00:00 2001
From: Jordan Dominion <Cyberboss@users.noreply.github.com>
Date: Mon, 26 Jun 2023 18:58:56 -0400
Subject: [PATCH] ValueTask IRestartHandler

Also fix a race condition with restarting
---
 .../Components/Chat/ChatManager.cs                   |  9 ++++-----
 .../Components/Watchdog/WatchdogBase.cs              |  2 +-
 src/Tgstation.Server.Host/Core/IRestartHandler.cs    |  4 ++--
 src/Tgstation.Server.Host/Server.cs                  | 12 ++++++++----
 4 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/Tgstation.Server.Host/Components/Chat/ChatManager.cs b/src/Tgstation.Server.Host/Components/Chat/ChatManager.cs
index 7b92808477..6690344fde 100644
--- a/src/Tgstation.Server.Host/Components/Chat/ChatManager.cs
+++ b/src/Tgstation.Server.Host/Components/Chat/ChatManager.cs
@@ -536,7 +536,7 @@ namespace Tgstation.Server.Host.Components.Chat
 		}
 
 		/// <inheritdoc />
-		public Task HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken)
+		public ValueTask HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken)
 		{
 			var message =
 				updateVersion == null
@@ -1001,7 +1001,7 @@ namespace Tgstation.Server.Host.Components.Chat
 		/// <param name="message">The <see cref="MessageContent"/> to send.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
 		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task SendMessage(IEnumerable<ulong> channelIds, Message replyTo, MessageContent message, CancellationToken cancellationToken)
+		ValueTask SendMessage(IEnumerable<ulong> channelIds, Message replyTo, MessageContent message, CancellationToken cancellationToken)
 		{
 			var channelIdsList = channelIds.ToList();
 
@@ -1012,7 +1012,7 @@ namespace Tgstation.Server.Host.Components.Chat
 				String.Join(", ", channelIdsList));
 
 			if (!channelIdsList.Any())
-				return Task.CompletedTask;
+				return ValueTask.CompletedTask;
 
 			return ValueTaskExtensions.WhenAll(
 				channelIdsList.Select(x =>
@@ -1027,8 +1027,7 @@ namespace Tgstation.Server.Host.Components.Chat
 							return ValueTask.CompletedTask;
 					return provider.SendMessage(replyTo, message, channelMapping.ProviderChannelId, cancellationToken);
 				}),
-				channelIdsList.Count)
-				.AsTask();
+				channelIdsList.Count);
 		}
 
 		/// <summary>
diff --git a/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs b/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
index a5eed5deaf..81a7c45d7a 100644
--- a/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
+++ b/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
@@ -402,7 +402,7 @@ namespace Tgstation.Server.Host.Components.Watchdog
 		}
 
 		/// <inheritdoc />
-		public async Task HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken)
+		public async ValueTask HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken)
 		{
 			if (gracefulShutdown)
 			{
diff --git a/src/Tgstation.Server.Host/Core/IRestartHandler.cs b/src/Tgstation.Server.Host/Core/IRestartHandler.cs
index 06549ff14f..db7d3e9e8c 100644
--- a/src/Tgstation.Server.Host/Core/IRestartHandler.cs
+++ b/src/Tgstation.Server.Host/Core/IRestartHandler.cs
@@ -15,7 +15,7 @@ namespace Tgstation.Server.Host.Core
 		/// <param name="updateVersion">The <see cref="Version"/> being updated to, <see langword="null"/> if not being changed.</param>
 		/// <param name="gracefulShutdown">If <see langword="true"/> the server should not expect to restart.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask HandleRestart(Version updateVersion, bool gracefulShutdown, CancellationToken cancellationToken);
 	}
 }
diff --git a/src/Tgstation.Server.Host/Server.cs b/src/Tgstation.Server.Host/Server.cs
index 6a26ed8440..2cbc26d2ec 100644
--- a/src/Tgstation.Server.Host/Server.cs
+++ b/src/Tgstation.Server.Host/Server.cs
@@ -10,6 +10,7 @@ using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
+using Tgstation.Server.Common.Extensions;
 using Tgstation.Server.Host.Configuration;
 using Tgstation.Server.Host.Core;
 
@@ -313,10 +314,13 @@ namespace Tgstation.Server.Host
 				var cancellationToken = cts.Token;
 				try
 				{
-					var eventsTask = Task.WhenAll(
-						restartHandlers.Select(
-							x => x.HandleRestart(newVersion, isGracefulShutdown, cancellationToken))
-						.ToList());
+					ValueTask eventsTask;
+					lock (restartLock)
+						eventsTask = ValueTaskExtensions.WhenAll(
+							restartHandlers
+								.Select(
+									x => x.HandleRestart(newVersion, isGracefulShutdown, cancellationToken))
+								.ToList());
 
 					logger.LogTrace("Joining restart handlers...");
 					await eventsTask;