From f38bc674055656f447d057ea506bf4183ad4f11a Mon Sep 17 00:00:00 2001
From: Tim Otten <totten@civicrm.org>
Date: Wed, 20 Jul 2022 01:11:22 -0700
Subject: [PATCH] Upgrader - Re-assess dispatch policy at the start of every
 upgrade task

The practical effect should be this:

* All tasks before `doCoreFinish()` run with the `upgrade.main` policy
* All tasks after `doCoreFinish()` run with the `upgrade.finish` policy

Framing this change is a bit delicate because tasks can run in different
ways (eg AJAX queue runner; headless queue-runner; cv queue-runner).  This
should have the same consequence in all runners.
---
 CRM/Upgrade/DispatchPolicy.php   | 26 ++++++++++++++++++++++++++
 CRM/Upgrade/Incremental/Base.php |  5 +----
 Civi/Core/Container.php          |  1 +
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/CRM/Upgrade/DispatchPolicy.php b/CRM/Upgrade/DispatchPolicy.php
index 087b70d88d..112f1797ad 100644
--- a/CRM/Upgrade/DispatchPolicy.php
+++ b/CRM/Upgrade/DispatchPolicy.php
@@ -150,4 +150,30 @@ class CRM_Upgrade_DispatchPolicy {
     return $policies[$phase];
   }
 
+  /**
+   * Assert that a specific policy is currently active.
+   *
+   * @param string $name
+   *   Ex: 'upgrade.main' or 'upgrade.finish'
+   * @throws \RuntimeException
+   */
+  public static function assertActive(string $name) {
+    $expected = static::get($name);
+    $actual = Civi::dispatcher()->getDispatchPolicy();
+    if ($expected != $actual) {
+      throw new \RuntimeException("Task can not execute correctly. The wrong dispatch policy is active. Expected to find \"$name\".");
+    }
+  }
+
+  /**
+   * Before running upgrade tasks, ensure that we apply the current dispatch-policy.
+   *
+   * @param \Civi\Core\Event\GenericHookEvent $event
+   */
+  public static function onRunTask(\Civi\Core\Event\GenericHookEvent $event) {
+    if ($event->taskCtx->queue->getName() === \CRM_Upgrade_Form::QUEUE_NAME) {
+      Civi::dispatcher()->setDispatchPolicy(static::pick());
+    }
+  }
+
 }
diff --git a/CRM/Upgrade/Incremental/Base.php b/CRM/Upgrade/Incremental/Base.php
index b109b5772f..e9bf28307d 100644
--- a/CRM/Upgrade/Incremental/Base.php
+++ b/CRM/Upgrade/Incremental/Base.php
@@ -213,10 +213,7 @@ class CRM_Upgrade_Incremental_Base {
     // The `enableExtension` has a very high value of `weight`, so this runs after all
     // core DB schema updates have been resolved. We can use high-level services.
 
-    Civi::dispatcher()->setDispatchPolicy(\CRM_Upgrade_DispatchPolicy::get('upgrade.finish'));
-    $restore = \CRM_Utils_AutoClean::with(function() {
-      Civi::dispatcher()->setDispatchPolicy(\CRM_Upgrade_DispatchPolicy::get('upgrade.main'));
-    });
+    CRM_Upgrade_DispatchPolicy::assertActive('upgrade.finish');
 
     $manager = CRM_Extension_System::singleton()->getManager();
     $manager->enable($manager->findInstallRequirements($keys));
diff --git a/Civi/Core/Container.php b/Civi/Core/Container.php
index f705463b7c..b33f950604 100644
--- a/Civi/Core/Container.php
+++ b/Civi/Core/Container.php
@@ -597,6 +597,7 @@ class Container {
     $bootServices['paths'] = new \Civi\Core\Paths();
 
     $bootServices['dispatcher.boot'] = new CiviEventDispatcher();
+    $bootServices['dispatcher.boot']->addListener('civi.queue.runTask.start', ['CRM_Upgrade_DispatchPolicy', 'onRunTask']);
 
     // Quality control: There should be no pre-boot hooks because they make it harder to understand/support/refactor.
     // If a pre-boot hook sneaks in, we'll raise an error.
-- 
2.25.1