From 3aafc02687d6ccb6c62193192df855b47051fc8f Mon Sep 17 00:00:00 2001
From: Qstick <qstick@gmail.com>
Date: Wed, 12 Aug 2020 22:28:57 -0400
Subject: [PATCH] Added update check early in startup if package requested

Fixes #4723
---
 .../Lifecycle/ApplicationStartingEvent.cs     |  8 +++
 .../Update/InstallUpdateService.cs            | 72 +++++++++++++++++--
 src/NzbDrone.Host/ApplicationServer.cs        |  8 +++
 3 files changed, 84 insertions(+), 4 deletions(-)
 create mode 100644 src/NzbDrone.Core/Lifecycle/ApplicationStartingEvent.cs

diff --git a/src/NzbDrone.Core/Lifecycle/ApplicationStartingEvent.cs b/src/NzbDrone.Core/Lifecycle/ApplicationStartingEvent.cs
new file mode 100644
index 000000000..464ff5ff4
--- /dev/null
+++ b/src/NzbDrone.Core/Lifecycle/ApplicationStartingEvent.cs
@@ -0,0 +1,8 @@
+using NzbDrone.Common.Messaging;
+
+namespace NzbDrone.Core.Lifecycle
+{
+    public class ApplicationStartingEvent : IEvent
+    {
+    }
+}
diff --git a/src/NzbDrone.Core/Update/InstallUpdateService.cs b/src/NzbDrone.Core/Update/InstallUpdateService.cs
index 4783f8074..96a761e96 100644
--- a/src/NzbDrone.Core/Update/InstallUpdateService.cs
+++ b/src/NzbDrone.Core/Update/InstallUpdateService.cs
@@ -1,6 +1,8 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
+using System.Threading;
 using NLog;
 using NzbDrone.Common;
 using NzbDrone.Common.Disk;
@@ -11,12 +13,14 @@
 using NzbDrone.Common.Processes;
 using NzbDrone.Core.Backup;
 using NzbDrone.Core.Configuration;
+using NzbDrone.Core.Lifecycle;
 using NzbDrone.Core.Messaging.Commands;
+using NzbDrone.Core.Messaging.Events;
 using NzbDrone.Core.Update.Commands;
 
 namespace NzbDrone.Core.Update
 {
-    public class InstallUpdateService : IExecute<ApplicationCheckUpdateCommand>, IExecute<ApplicationUpdateCommand>
+    public class InstallUpdateService : IExecute<ApplicationCheckUpdateCommand>, IExecute<ApplicationUpdateCommand>, IHandle<ApplicationStartingEvent>
     {
         private readonly ICheckUpdateService _checkUpdateService;
         private readonly Logger _logger;
@@ -75,7 +79,7 @@ public InstallUpdateService(ICheckUpdateService checkUpdateService,
             _logger = logger;
         }
 
-        private void InstallUpdate(UpdatePackage updatePackage)
+        private bool InstallUpdate(UpdatePackage updatePackage)
         {
             EnsureAppDataSafety();
 
@@ -98,7 +102,7 @@ private void InstallUpdate(UpdatePackage updatePackage)
             if (_appFolderInfo.StartUpFolder.EndsWith("_output"))
             {
                 _logger.ProgressDebug("Running in developer environment, not updating.");
-                return;
+                return false;
             }
 
             var updateSandboxFolder = _appFolderInfo.GetUpdateSandboxFolder();
@@ -136,7 +140,7 @@ private void InstallUpdate(UpdatePackage updatePackage)
             if (OsInfo.IsNotWindows && _configFileProvider.UpdateMechanism == UpdateMechanism.Script)
             {
                 InstallUpdateWithScript(updateSandboxFolder);
-                return;
+                return true;
             }
 
             _logger.Info("Preparing client");
@@ -152,6 +156,8 @@ private void InstallUpdate(UpdatePackage updatePackage)
             _logger.ProgressInfo("Radarr will restart shortly.");
 
             _processProvider.Start(_appFolderInfo.GetUpdateClientExePath(updatePackage.Runtime), GetUpdaterArgs(updateSandboxFolder));
+
+            return true;
         }
 
         private void EnsureValidBranch(UpdatePackage package)
@@ -286,5 +292,63 @@ public void Execute(ApplicationUpdateCommand message)
                 }
             }
         }
+
+        public void Handle(ApplicationStartingEvent message)
+        {
+            // Check if we have to do an application update on startup
+            try
+            {
+                // Don't do a prestartup update check unless BuiltIn update is enabled
+                if (_configFileProvider.UpdateAutomatically ||
+                    _configFileProvider.UpdateMechanism != UpdateMechanism.BuiltIn ||
+                    _deploymentInfoProvider.IsExternalUpdateMechanism)
+                {
+                    return;
+                }
+
+                var updateMarker = Path.Combine(_appFolderInfo.AppDataFolder, "update_required");
+                if (!_diskProvider.FileExists(updateMarker))
+                {
+                    return;
+                }
+
+                _logger.Debug("Post-install update check requested");
+
+                var latestAvailable = _checkUpdateService.AvailableUpdate();
+                if (latestAvailable == null)
+                {
+                    _logger.Debug("No post-install update available");
+                    _diskProvider.DeleteFile(updateMarker);
+                    return;
+                }
+
+                _logger.Info("Installing post-install update from {0} to {1}", BuildInfo.Version, latestAvailable.Version);
+                _diskProvider.DeleteFile(updateMarker);
+
+                var installing = InstallUpdate(latestAvailable);
+
+                if (installing)
+                {
+                    _logger.Debug("Install in progress, giving installer 30 seconds.");
+
+                    var watch = Stopwatch.StartNew();
+
+                    while (watch.Elapsed < TimeSpan.FromSeconds(30))
+                    {
+                        Thread.Sleep(1000);
+                    }
+
+                    _logger.Error("Post-install update not completed within 30 seconds. Attempting to continue normal operation.");
+                }
+                else
+                {
+                    _logger.Debug("Post-install update cancelled for unknown reason. Attempting to continue normal operation.");
+                }
+            }
+            catch (Exception ex)
+            {
+                _logger.Error(ex, "Failed to perform the post-install update check. Attempting to continue normal operation.");
+            }
+        }
     }
 }
diff --git a/src/NzbDrone.Host/ApplicationServer.cs b/src/NzbDrone.Host/ApplicationServer.cs
index 2e1bdf961..7d6ae034b 100644
--- a/src/NzbDrone.Host/ApplicationServer.cs
+++ b/src/NzbDrone.Host/ApplicationServer.cs
@@ -91,6 +91,14 @@ public void Start()
 
             _runtimeInfo.IsExiting = false;
             DbFactory.RegisterDatabase(_container);
+
+            _container.Resolve<IEventAggregator>().PublishEvent(new ApplicationStartingEvent());
+
+            if (_runtimeInfo.IsExiting)
+            {
+                return;
+            }
+
             _hostController.StartServer();
 
             if (!_startupContext.Flags.Contains(StartupContext.NO_BROWSER)