From 362adcab1acf4bc4fa928f0184a90c147756e1c8 Mon Sep 17 00:00:00 2001
From: Yonas Habteab <yonas.habteab@icinga.com>
Date: Mon, 7 Mar 2022 17:50:15 +0100
Subject: [PATCH] ConfigPackageUtility: Don't reset ongoing package updates on
 config validation success and process is going to be reloaded

---
 lib/remote/configpackageutility.cpp | 21 +++++++++++++++++----
 lib/remote/configpackageutility.hpp |  8 ++++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/lib/remote/configpackageutility.cpp b/lib/remote/configpackageutility.cpp
index b70009997..e79540147 100644
--- a/lib/remote/configpackageutility.cpp
+++ b/lib/remote/configpackageutility.cpp
@@ -167,7 +167,8 @@ void ConfigPackageUtility::ActivateStage(const String& packageName, const String
 	WritePackageConfig(packageName);
 }
 
-void ConfigPackageUtility::TryActivateStageCallback(const ProcessResult& pr, const String& packageName, const String& stageName, bool activate, bool reload)
+void ConfigPackageUtility::TryActivateStageCallback(const ProcessResult& pr, const String& packageName, const String& stageName,
+	bool activate, bool reload, const Shared<Defer>::Ptr& resetPackageUpdates)
 {
 	String logFile = GetPackageDir() + "/" + packageName + "/" + stageName + "/startup.log";
 	std::ofstream fpLog(logFile.CStr(), std::ofstream::out | std::ostream::binary | std::ostream::trunc);
@@ -188,8 +189,17 @@ void ConfigPackageUtility::TryActivateStageCallback(const ProcessResult& pr, con
 				ActivateStage(packageName, stageName);
 			}
 
-			if (reload)
+			if (reload) {
+				/*
+				 * Cancel the deferred callback before going out of scope so that the config stages handler
+				 * flag isn't resetting earlier and allowing other clients to submit further requests while
+				 * Icinga2 is reloading. Otherwise, the ongoing request will be cancelled halfway before the
+				 * operation is completed once the new worker becomes ready.
+				 */
+				resetPackageUpdates->Cancel();
+
 				Application::RequestRestart();
+			}
 		}
 	} else {
 		Log(LogCritical, "ConfigPackageUtility")
@@ -198,7 +208,8 @@ void ConfigPackageUtility::TryActivateStageCallback(const ProcessResult& pr, con
 	}
 }
 
-void ConfigPackageUtility::AsyncTryActivateStage(const String& packageName, const String& stageName, bool activate, bool reload)
+void ConfigPackageUtility::AsyncTryActivateStage(const String& packageName, const String& stageName, bool activate, bool reload,
+	const Shared<Defer>::Ptr& resetPackageUpdates)
 {
 	VERIFY(Application::GetArgC() >= 1);
 
@@ -224,7 +235,9 @@ void ConfigPackageUtility::AsyncTryActivateStage(const String& packageName, cons
 
 	Process::Ptr process = new Process(Process::PrepareCommand(args));
 	process->SetTimeout(Application::GetReloadTimeout());
-	process->Run([packageName, stageName, activate, reload](const ProcessResult& pr) { TryActivateStageCallback(pr, packageName, stageName, activate, reload); });
+	process->Run([packageName, stageName, activate, reload, resetPackageUpdates](const ProcessResult& pr) {
+		TryActivateStageCallback(pr, packageName, stageName, activate, reload, resetPackageUpdates);
+	});
 }
 
 void ConfigPackageUtility::DeleteStage(const String& packageName, const String& stageName)
diff --git a/lib/remote/configpackageutility.hpp b/lib/remote/configpackageutility.hpp
index 5c16b75ee..240f5919c 100644
--- a/lib/remote/configpackageutility.hpp
+++ b/lib/remote/configpackageutility.hpp
@@ -8,6 +8,8 @@
 #include "base/dictionary.hpp"
 #include "base/process.hpp"
 #include "base/string.hpp"
+#include "base/defer.hpp"
+#include "base/shared.hpp"
 #include <vector>
 
 namespace icinga
@@ -37,7 +39,8 @@ public:
 	static void SetActiveStage(const String& packageName, const String& stageName);
 	static void SetActiveStageToFile(const String& packageName, const String& stageName);
 	static void ActivateStage(const String& packageName, const String& stageName);
-	static void AsyncTryActivateStage(const String& packageName, const String& stageName, bool activate, bool reload);
+	static void AsyncTryActivateStage(const String& packageName, const String& stageName, bool activate, bool reload,
+		const Shared<Defer>::Ptr& resetPackageUpdates);
 
 	static std::vector<std::pair<String, bool> > GetFiles(const String& packageName, const String& stageName);
 
@@ -59,7 +62,8 @@ private:
 	static void WritePackageConfig(const String& packageName);
 	static void WriteStageConfig(const String& packageName, const String& stageName);
 
-	static void TryActivateStageCallback(const ProcessResult& pr, const String& packageName, const String& stageName, bool activate, bool reload);
+	static void TryActivateStageCallback(const ProcessResult& pr, const String& packageName, const String& stageName, bool activate,
+		bool reload, const Shared<Defer>::Ptr& resetPackageUpdates);
 
 	static bool ValidateFreshName(const String& name);
 };