Workflow engine graceful shutdown

Author: Khushboo

st2actions/st2actions/workflows/workflows.py

@@ -14,22 +14,29 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+from oslo_config import cfg
 
 from orquesta import statuses
-
+from tooz.coordination import GroupNotCreated
+from st2common.services import coordination
+from eventlet.semaphore import Semaphore
+from eventlet import spawn_after
 from st2common.constants import action as ac_const
 from st2common import log as logging
 from st2common.metrics import base as metrics
 from st2common.models.db import execution as ex_db_models
 from st2common.models.db import workflow as wf_db_models
 from st2common.persistence import liveaction as lv_db_access
 from st2common.persistence import workflow as wf_db_access
+from st2common.persistence import execution as ex_db_access
+from st2common.services import action as ac_svc
 from st2common.services import policies as pc_svc
 from st2common.services import workflows as wf_svc
 from st2common.transport import consumers
 from st2common.transport import queues
 from st2common.transport import utils as txpt_utils
-
+from st2common.util import concurrency
+from st2common.util import action_db as action_utils
 
 LOG = logging.getLogger(__name__)
 
@@ -40,10 +47,15 @@
     queues.WORKFLOW_ACTION_EXECUTION_UPDATE_QUEUE,
 ]
 
+WORKFLOW_ENGINE = "workflow_engine"
+SHUTDOWN_ROUTINE = "shutdown_routine"
+
 
 class WorkflowExecutionHandler(consumers.VariableMessageHandler):
     def __init__(self, connection, queues):
         super(WorkflowExecutionHandler, self).__init__(connection, queues)
+        self._active_messages = 0
+        self._semaphore = Semaphore()
 
         def handle_workflow_execution_with_instrumentation(wf_ex_db):
             with metrics.CounterWithTimer(key="orquesta.workflow.executions"):
@@ -62,6 +74,10 @@ def handle_action_execution_with_instrumentation(ac_ex_db):
             ex_db_models.ActionExecutionDB: handle_action_execution_with_instrumentation,
         }
 
+        # This is required to ensure workflows stuck in pausing state after shutdown transition to paused state after engine startup.
+        self._delay = 30
+        spawn_after(self._delay, self._resume_workflows_paused_during_shutdown)
+
     def get_queue_consumer(self, connection, queues):
         # We want to use a special ActionsQueueConsumer which uses 2 dispatcher pools
         return consumers.VariableMessageQueueConsumer(
@@ -78,13 +94,61 @@ def process(self, message):
             raise ValueError(msg)
 
         try:
+            with self._semaphore:
+                self._active_messages += 1
             handler_function(message)
         except Exception as e:
             # If the exception is caused by DB connection error, then the following
             # error handling routine will fail as well because it will try to update
             # the database and fail the workflow execution gracefully. In this case,
             # the garbage collector will find and cancel these workflow executions.
             self.fail_workflow_execution(message, e)
+        finally:
+            with self._semaphore:
+                self._active_messages -= 1
+
+    def shutdown(self):
+        super(WorkflowExecutionHandler, self).shutdown()
+        while self._active_messages > 0:
+            concurrency.sleep(2)
+
+        coordinator = coordination.get_coordinator()
+        member_ids = []
+        with coordinator.get_lock(SHUTDOWN_ROUTINE):
+            try:
+                member_ids = list(
+                    coordinator.get_members(WORKFLOW_ENGINE.encode("utf-8")).get()
+                )
+            except GroupNotCreated:
+                pass
+
+            # Check if there are other runners in service registry
+            if cfg.CONF.coordination.service_registry and not member_ids:
+                ac_ex_dbs = self._get_running_workflows()
+                for ac_ex_db in ac_ex_dbs:
+                    lv_ac = action_utils.get_liveaction_by_id(ac_ex_db.liveaction["id"])
+                    ac_svc.request_pause(lv_ac, SHUTDOWN_ROUTINE)
+
+    def _get_running_workflows(self):
+        query_filters = {
+            "runner__name": "orquesta",
+            "status": ac_const.LIVEACTION_STATUS_RUNNING,
+        }
+        return ex_db_access.ActionExecution.query(**query_filters)
+
+    def _get_workflows_paused_during_shutdown(self):
+        query_filters = {
+            "status": ac_const.LIVEACTION_STATUS_PAUSED,
+            "context__paused_by": SHUTDOWN_ROUTINE,
+        }
+        return lv_db_access.LiveAction.query(**query_filters)
+
+    def _resume_workflows_paused_during_shutdown(self):
+        coordinator = coordination.get_coordinator()
+        with coordinator.get_lock(SHUTDOWN_ROUTINE):
+            lv_ac_dbs = self._get_workflows_paused_during_shutdown()
+            for lv_ac_db in lv_ac_dbs:
+                ac_svc.request_resume(lv_ac_db, SHUTDOWN_ROUTINE)
 
     def fail_workflow_execution(self, message, exception):
         # Prepare attributes based on message type.

st2actions/tests/unit/test_workflow_engine.py

@@ -271,3 +271,60 @@ def test_process_error_handling_has_error(self, mock_get_lock):
         # Assert workflow execution is cleaned up and canceled.
         lv_ac_db = lv_db_access.LiveAction.get_by_id(str(lv_ac_db.id))
         self.assertEqual(lv_ac_db.status, action_constants.LIVEACTION_STATUS_CANCELED)
+
+    def test_workflow_engine_shutdown(self):
+        cfg.CONF.set_override(
+            name="service_registry", override=True, group="coordination"
+        )
+        wf_meta = self.get_wf_fixture_meta_data(TEST_PACK_PATH, "sequential.yaml")
+        lv_ac_db = lv_db_models.LiveActionDB(action=wf_meta["name"])
+        lv_ac_db, ac_ex_db = action_service.request(lv_ac_db)
+
+        # Assert action execution is running.
+        lv_ac_db = lv_db_access.LiveAction.get_by_id(str(lv_ac_db.id))
+        self.assertEqual(lv_ac_db.status, action_constants.LIVEACTION_STATUS_RUNNING)
+        wf_ex_db = wf_db_access.WorkflowExecution.query(
+            action_execution=str(ac_ex_db.id)
+        )[0]
+        self.assertEqual(wf_ex_db.status, action_constants.LIVEACTION_STATUS_RUNNING)
+        workflow_engine = workflows.get_engine()
+
+        # Manually add running workflow
+        workflow_engine._handling_workflows = [str(ac_ex_db.id)]
+        eventlet.spawn(workflow_engine.shutdown)
+
+        # Sleep for few seconds to ensure execution transitions to pausing.
+        eventlet.sleep(5)
+
+        lv_ac_db = lv_db_access.LiveAction.get_by_id(str(lv_ac_db.id))
+        self.assertEqual(lv_ac_db.status, action_constants.LIVEACTION_STATUS_PAUSING)
+
+        # Process task1.
+        query_filters = {"workflow_execution": str(wf_ex_db.id), "task_id": "task1"}
+        t1_ex_db = wf_db_access.TaskExecution.query(**query_filters)[0]
+        t1_ac_ex_db = ex_db_access.ActionExecution.query(
+            task_execution=str(t1_ex_db.id)
+        )[0]
+
+        workflows.get_engine().process(t1_ac_ex_db)
+        t1_ac_ex_db = ex_db_access.ActionExecution.query(
+            task_execution=str(t1_ex_db.id)
+        )[0]
+        self.assertEqual(
+            t1_ac_ex_db.status, action_constants.LIVEACTION_STATUS_SUCCEEDED
+        )
+
+        lv_ac_db = lv_db_access.LiveAction.get_by_id(str(lv_ac_db.id))
+        self.assertEqual(lv_ac_db.status, action_constants.LIVEACTION_STATUS_PAUSED)
+
+        workflow_engine = workflows.get_engine()
+        eventlet.sleep(workflow_engine._delay)
+        lv_ac_db = lv_db_access.LiveAction.get_by_id(str(lv_ac_db.id))
+        self.assertTrue(
+            lv_ac_db.status
+            in [
+                action_constants.LIVEACTION_STATUS_RESUMING,
+                action_constants.LIVEACTION_STATUS_RUNNING,
+                action_constants.LIVEACTION_STATUS_SUCCEEDED,
+            ]
+        )

st2common/st2common/transport/consumers.py

@@ -43,6 +43,7 @@ def __init__(self, connection, queues, handler):
         self._handler = handler
 
     def shutdown(self):
+        self.should_stop = True
         self._dispatcher.shutdown()
 
     def get_consumers(self, Consumer, channel):