From 709b2c514eb15b9e73ab506293c5c519dd3fffca Mon Sep 17 00:00:00 2001
From: Dominion <Cyberboss@users.noreply.github.com>
Date: Sun, 11 Jun 2023 11:19:18 -0400
Subject: [PATCH] Trying to track down elusive swarm update errors

---
 .../Core/ServerUpdater.cs                     |  6 ++++-
 .../IO/BufferedFileStreamProvider.cs          | 26 ++++++++++---------
 .../Swarm/SwarmService.cs                     | 16 +++++++++---
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/Tgstation.Server.Host/Core/ServerUpdater.cs b/src/Tgstation.Server.Host/Core/ServerUpdater.cs
index 07f98fa7d3..3de6a6088c 100644
--- a/src/Tgstation.Server.Host/Core/ServerUpdater.cs
+++ b/src/Tgstation.Server.Host/Core/ServerUpdater.cs
@@ -125,6 +125,7 @@ namespace Tgstation.Server.Host.Core
 					return false;
 
 				await using var bufferedStream = tuple.Item1;
+				var needStreamUntilCommit = tuple.Item2;
 				var createdStagingDirectory = false;
 				try
 				{
@@ -136,8 +137,11 @@ namespace Tgstation.Server.Host.Core
 						createdStagingDirectory = true;
 						await ioManager.ZipToDirectory(stagingDirectory, updateZipData, cancellationToken);
 
-						if (!tuple.Item2)
+						if (!needStreamUntilCommit)
+						{
+							logger.LogTrace("Early disposing update stream provider...");
 							await bufferedStream.DisposeAsync(); // don't leave this in memory
+						}
 					}
 					catch (Exception ex)
 					{
diff --git a/src/Tgstation.Server.Host/IO/BufferedFileStreamProvider.cs b/src/Tgstation.Server.Host/IO/BufferedFileStreamProvider.cs
index a92e28f998..3542266176 100644
--- a/src/Tgstation.Server.Host/IO/BufferedFileStreamProvider.cs
+++ b/src/Tgstation.Server.Host/IO/BufferedFileStreamProvider.cs
@@ -62,7 +62,7 @@ namespace Tgstation.Server.Host.IO
 			lock (semaphore)
 			{
 				localBuffer = buffer;
-				if (buffered && localBuffer == null)
+				if (localBuffer == null)
 					return;
 
 				// important to drop the reference so it can properly GC
@@ -101,21 +101,23 @@ namespace Tgstation.Server.Host.IO
 				using (await SemaphoreSlimContext.Lock(semaphore, cancellationToken))
 					if (!buffered)
 					{
-						await input.CopyToAsync(buffer, cancellationToken);
+						MemoryStream localBuffer;
 						lock (semaphore)
-						{
-							if (buffer == null)
-								throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
+							localBuffer = buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
 
-							buffer.Seek(0, SeekOrigin.Begin);
-							buffered = true;
-							return (buffer, buffer.Length);
-						}
+						await input.CopyToAsync(localBuffer, cancellationToken);
+						localBuffer.Seek(0, SeekOrigin.Begin);
+						buffered = true;
+						return (localBuffer, localBuffer.Length);
 					}
 
-			return (
-				buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider)),
-				buffer.Length);
+			lock (semaphore)
+			{
+				var localBuffer = buffer ?? throw new ObjectDisposedException(nameof(BufferedFileStreamProvider));
+				return (
+					localBuffer,
+					localBuffer.Length);
+			}
 		}
 	}
 }
diff --git a/src/Tgstation.Server.Host/Swarm/SwarmService.cs b/src/Tgstation.Server.Host/Swarm/SwarmService.cs
index 539cb75477..b3c5cec6ac 100644
--- a/src/Tgstation.Server.Host/Swarm/SwarmService.cs
+++ b/src/Tgstation.Server.Host/Swarm/SwarmService.cs
@@ -811,7 +811,7 @@ namespace Tgstation.Server.Host.Swarm
 
 				if (!swarmController && initiator)
 				{
-					var downloadTickets = CreateDownloadTickets(initiatorProvider, currentNodes);
+					var downloadTickets = await CreateDownloadTickets(initiatorProvider, currentNodes, cancellationToken);
 
 					logger.LogInformation("Forwarding update request to swarm controller...");
 					using var httpClient = httpClientFactory.CreateClient();
@@ -964,7 +964,7 @@ namespace Tgstation.Server.Host.Swarm
 				}
 
 				var downloadTicketDictionary = weAreInitiator
-					? CreateDownloadTickets(initiatorProvider, currentUpdateOperation.InvolvedServers)
+					? await CreateDownloadTickets(initiatorProvider, currentUpdateOperation.InvolvedServers, cancellationToken)
 					: updateRequest.DownloadTickets;
 
 				var sourceNode = weAreInitiator
@@ -1058,9 +1058,16 @@ namespace Tgstation.Server.Host.Swarm
 		/// </summary>
 		/// <param name="initiatorProvider">The <see cref="ISeekableFileStreamProvider"/> containing the server update package.</param>
 		/// <param name="involvedServers">An <see cref="IEnumerable{T}"/> of the involved <see cref="SwarmServerResponse"/>.</param>
-		/// <returns>A new <see cref="Dictionary{TKey, TValue}"/> of unique <see cref="FileTicketResponse"/>s keyed by their <see cref="Api.Models.Internal.SwarmServer.Identifier"/>.</returns>
-		Dictionary<string, FileTicketResponse> CreateDownloadTickets(ISeekableFileStreamProvider initiatorProvider, IReadOnlyCollection<SwarmServerResponse> involvedServers)
+		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
+		/// <returns>A <see cref="Task{TResult}"/> resulting in a new <see cref="Dictionary{TKey, TValue}"/> of unique <see cref="FileTicketResponse"/>s keyed by their <see cref="Api.Models.Internal.SwarmServer.Identifier"/>.</returns>
+		async Task<Dictionary<string, FileTicketResponse>> CreateDownloadTickets(
+			ISeekableFileStreamProvider initiatorProvider,
+			IReadOnlyCollection<SwarmServerResponse> involvedServers,
+			CancellationToken cancellationToken)
 		{
+			// we need to ensure this thing is loaded before we start providing downloads or it'll create unnecessary delays
+			var bufferingTask = initiatorProvider.GetResult(cancellationToken);
+
 			var downloadProvider = new FileDownloadProvider(
 				() => initiatorProvider.Disposed
 					? Api.Models.ErrorCode.ResourceNotPresent
@@ -1081,6 +1088,7 @@ namespace Tgstation.Server.Host.Swarm
 					node.Identifier,
 					transferService.CreateDownload(downloadProvider));
 
+			await bufferingTask;
 			return downloadTickets;
 		}