From 9d03e36b0dc56e2cac9c20cf1529c1b53cea73b8 Mon Sep 17 00:00:00 2001
From: Jordan Brown <Cyberboss@users.noreply.github.com>
Date: Wed, 8 Jul 2020 23:26:34 -0400
Subject: [PATCH] Improve session persistor slightly

---
 .../Components/Session/SessionPersistor.cs    | 29 ++++++++++---------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/src/Tgstation.Server.Host/Components/Session/SessionPersistor.cs b/src/Tgstation.Server.Host/Components/Session/SessionPersistor.cs
index 90330df437..3e5dce2043 100644
--- a/src/Tgstation.Server.Host/Components/Session/SessionPersistor.cs
+++ b/src/Tgstation.Server.Host/Components/Session/SessionPersistor.cs
@@ -76,11 +76,14 @@ namespace Tgstation.Server.Host.Components.Session
 				.DeleteAsync(cancellationToken)
 				.ConfigureAwait(false);
 
-			db.CompileJobs.Attach(reattachInformation.Dmb.CompileJob);
+			var compileJob = reattachInformation.Dmb.CompileJob;
+			db.Instances.Attach(compileJob.Job.Instance);
+			db.Jobs.Add(compileJob.Job);
+			db.CompileJobs.Attach(compileJob);
 			var dbReattachInfo = new Models.ReattachInformation
 			{
 				AccessIdentifier = reattachInformation.AccessIdentifier,
-				CompileJob = reattachInformation.Dmb.CompileJob,
+				CompileJob = compileJob,
 				Port = reattachInformation.Port,
 				ProcessId = reattachInformation.ProcessId,
 				RebootState = reattachInformation.RebootState,
@@ -98,6 +101,16 @@ namespace Tgstation.Server.Host.Components.Session
 			TimeSpan? topicTimeout = null;
 			await databaseContextFactory.UseContext(async (db) =>
 			{
+				var dbReattachInfos = await db
+					.ReattachInformations
+					.AsQueryable()
+					.Where(x => x.CompileJob.Job.Instance.Id == metadata.Id)
+					.Include(x => x.CompileJob)
+					.ToListAsync(cancellationToken).ConfigureAwait(false);
+				result = dbReattachInfos.FirstOrDefault();
+				if (result == default)
+					return;
+
 				var timeoutMilliseconds = await db
 					.Instances
 					.AsQueryable()
@@ -114,16 +127,6 @@ namespace Tgstation.Server.Host.Components.Session
 
 				topicTimeout = TimeSpan.FromMilliseconds(timeoutMilliseconds.Value);
 
-				var dbReattachInfos = await db
-					.ReattachInformations
-					.AsQueryable()
-					.Where(x => x.CompileJob.Job.Instance.Id == metadata.Id)
-					.Include(x => x.CompileJob)
-					.ToListAsync(cancellationToken).ConfigureAwait(false);
-				result = dbReattachInfos.FirstOrDefault();
-				if (result == default)
-					return;
-
 				bool first = true;
 				foreach (var reattachInfo in dbReattachInfos)
 				{
@@ -148,7 +151,7 @@ namespace Tgstation.Server.Host.Components.Session
 				await db.Save(cancellationToken).ConfigureAwait(false);
 			}).ConfigureAwait(false);
 
-			if (result == default)
+			if (!topicTimeout.HasValue)
 			{
 				logger.LogDebug("Reattach information not found!");
 				return null;