From 5aa4645a9315277991a5aa971674dfb89446967a Mon Sep 17 00:00:00 2001
From: Madhu Kanoor <mkanoor@redhat.com>
Date: Thu, 26 Jan 2023 14:01:07 -0500
Subject: [PATCH] Shutdown can optionally include message and delay

https://issues.redhat.com/browse/AAP-8646

The Shutdown can optionally take the following parameters
- message
- delay

When a source ends it puts out a message
```
Source range initiated shutdown at 2023-01-26 16:35:49.250861
```
---
 ansible_rulebook/builtin.py                 | 11 +++++++
 ansible_rulebook/engine.py                  | 36 +++++++++++++++++----
 ansible_rulebook/messages.py                |  5 +--
 ansible_rulebook/schema/ruleset_schema.json |  7 +++-
 docs/actions.rst                            | 17 +++++++++-
 tests/examples/38_shutdown.yml              |  2 ++
 tests/test_examples.py                      |  4 +++
 7 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/ansible_rulebook/builtin.py b/ansible_rulebook/builtin.py
index ad7a47b1..457acf68 100644
--- a/ansible_rulebook/builtin.py
+++ b/ansible_rulebook/builtin.py
@@ -776,6 +776,8 @@ async def shutdown(
     source_ruleset_name: str,
     source_rule_name: str,
     ruleset: str,
+    delay: float = 0.0,
+    message: str = "Default shutdown message",
 ):
     await event_log.put(
         dict(
@@ -786,8 +788,17 @@ async def shutdown(
             rule=source_rule_name,
             run_at=str(datetime.utcnow()),
             matching_events=_get_events(variables),
+            delay=delay,
+            message=message,
         )
     )
+
+    print(
+        "Ruleset: %s rule: %s has initiated shutdown. "
+        "Delay: %.3f seconds, Message: %s"
+        % (source_ruleset_name, source_rule_name, delay, message)
+    )
+    await asyncio.sleep(delay)
     raise ShutdownException()
 
 
diff --git a/ansible_rulebook/engine.py b/ansible_rulebook/engine.py
index 0bbb34f7..265b9628 100644
--- a/ansible_rulebook/engine.py
+++ b/ansible_rulebook/engine.py
@@ -156,16 +156,32 @@ async def start_source(
             raise Exception("Entrypoint is not a coroutine function.")
 
         await entrypoint(fqueue, args)
+        shutdown_msg = (
+            f"Source {source.source_name} initiated shutdown at "
+            f"{str(datetime.now())}"
+        )
 
     except KeyboardInterrupt:
+        shutdown_msg = (
+            f"Source {source.source_name} keyboard interrupt, "
+            f"initiated shutdown at {str(datetime.now())}"
+        )
         pass
     except asyncio.CancelledError:
-        logger.info("Task cancelled")
-    except BaseException:
+        shutdown_msg = (
+            f"Source {source.source_name} task cancelled, "
+            "initiated shutdown at {str(datetime.now())}"
+        )
+        logger.info("Task cancelled " + shutdown_msg)
+    except BaseException as e:
         logger.exception("Source error")
+        shutdown_msg = (
+            f"Shutting down source: {source.source_name} error : {e}"
+        )
+        logger.error(shutdown_msg)
         raise
     finally:
-        await queue.put(Shutdown())
+        await queue.put(Shutdown(shutdown_msg, 0.0, source.source_name))
 
 
 async def run_rulesets(
@@ -260,7 +276,7 @@ async def run_ruleset(self):
         source_loop_task = asyncio.create_task(self._drain_source_queue())
         await asyncio.wait([source_loop_task])
 
-    async def _stop(self):
+    async def _stop(self, obj):
         # Wait for items in queues to be completed. Mainly for spec tests.
         await asyncio.sleep(0.01)
 
@@ -274,7 +290,14 @@ async def _stop(self):
             pass
 
         await asyncio.wait([self.pa_runner_task, self.action_loop_task])
-        await self.event_log.put(dict(type="Shutdown"))
+        await self.event_log.put(
+            dict(
+                type="Shutdown",
+                message=obj.message,
+                delay=obj.delay,
+                source_plugin=obj.source_plugin,
+            )
+        )
         lang.end_session(self.name)
 
     async def _drain_source_queue(self):
@@ -287,7 +310,8 @@ async def _drain_source_queue(self):
             logger.debug("Received event : " + str(data))
             json_count(data)
             if isinstance(data, Shutdown):
-                await self._stop()
+                logger.info("Shutdown message received: " + str(data))
+                await self._stop(data)
                 logger.info("Stopped waiting on events from %s", self.name)
                 return
             if not data:
diff --git a/ansible_rulebook/messages.py b/ansible_rulebook/messages.py
index c7bc817e..04803c51 100644
--- a/ansible_rulebook/messages.py
+++ b/ansible_rulebook/messages.py
@@ -16,5 +16,6 @@
 
 
 class Shutdown(NamedTuple):
-
-    pass
+    message: str = "Not specified"
+    delay: float = 0.0
+    source_plugin: str = ""
diff --git a/ansible_rulebook/schema/ruleset_schema.json b/ansible_rulebook/schema/ruleset_schema.json
index 8ffc0cb5..3a64042e 100644
--- a/ansible_rulebook/schema/ruleset_schema.json
+++ b/ansible_rulebook/schema/ruleset_schema.json
@@ -382,7 +382,12 @@
             "type": "object",
             "properties": {
                 "shutdown": {
-                    "type": ["object","null"]
+                    "type": ["object","null"],
+                    "properties": {
+                        "delay": {"type": "number" },
+                        "message": {"type": "string" }
+                    },
+                    "additionalProperties": false
                 }
             },
             "required": [
diff --git a/docs/actions.rst b/docs/actions.rst
index 81ee21d9..1f47f7a5 100644
--- a/docs/actions.rst
+++ b/docs/actions.rst
@@ -314,8 +314,21 @@ Example with multiple event match:
 
 shutdown
 ********
+.. list-table:: Shutdown ansible-rulebook
+   :widths: 25 150 10
+   :header-rows: 1
+
+   * - Name
+     - Description
+     - Required
+   * - delay
+     - A numeric value about how long to wait before shutting down, default 0.0
+     - No
+   * - message
+     - A message to be associated with this shutdown 
+     - No
 
-| Generate a shutdown event which will terminate the rulebook engine. If there are multiple
+| Generate a shutdown event which will terminate the ansible-rulebook process.
 | If there are multiple rule-sets running in your rule book, issuing a shutdown will cause
 | all other rule-sets to end, care needs to be taken to account for running playbooks which
 | can be impacted when one of the rule set decides to shutdown.
@@ -328,6 +341,8 @@ Example:
        condition: event.i >= 5
        action:
           shutdown:
+            delay: 0.125
+            message: Shutting down after 5 events
 
 Results
 *******
diff --git a/tests/examples/38_shutdown.yml b/tests/examples/38_shutdown.yml
index 49709927..e8daeff1 100644
--- a/tests/examples/38_shutdown.yml
+++ b/tests/examples/38_shutdown.yml
@@ -10,3 +10,5 @@
       condition: event.i == 1
       action:
         shutdown:
+          delay: 1.1845
+          message: My rule has triggered a shutdown
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 089d2970..cd85162e 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -990,6 +990,10 @@ async def test_38_shutdown_action():
     event = event_log.get_nowait()
     assert event["type"] == "Action", "1"
     assert event["action"] == "shutdown", "1"
+    assert event["message"] == "My rule has triggered a shutdown"
+    assert event["delay"] == 1.1845
+    assert event["ruleset"] == "Test shutdown action"
+    assert event["rule"] == "Host 1 rule"
     event = event_log.get_nowait()
     assert event["type"] == "Shutdown", "2"