From 159bd35d2f794b93e13ed73a74706d5e893e4b69 Mon Sep 17 00:00:00 2001
From: Dinesh Yadav <dy270k@att.com>
Date: Tue, 19 Jan 2021 22:42:13 +0000
Subject: [PATCH] Enhance Shaker to support the shell class execution after
 reboot agent instance.

Issue: Some of our complex test requires to reboot compute nodes. In
that scenario, if agent instances are on those computes then those
instances will also be rebooted and running shaker test(s) will be
stopped.

At this time this commit is specifically for shaker's shell class
execution as it writes execution script in /tmp directory and we
know that if system/instance is being rebooted then whatever /tmp
directory contains that will be wiped out. So, if user has a choice
to copy that execution script in some specific directory which is
not being impacted (i.e /mnt, /opt etc) after reboot then through
custom test scripts above condition can be handled

Change-Id: I41f824561e36329d07afc9e798eab44b72aeb01a
---
 doc/source/tools/shaker-agent.txt      |  6 +++++-
 doc/source/tools/shaker-all-in-one.txt |  7 ++++++-
 doc/source/tools/shaker.txt            |  7 ++++++-
 etc/shaker.conf                        |  5 +++++
 shaker/agent/agent.py                  | 10 +++++++++-
 shaker/engine/config.py                |  4 ++++
 6 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/doc/source/tools/shaker-agent.txt b/doc/source/tools/shaker-agent.txt
index f79a210..a25b1f9 100644
--- a/doc/source/tools/shaker-agent.txt
+++ b/doc/source/tools/shaker-agent.txt
@@ -1,4 +1,4 @@
-usage: shaker-agent [-h] [--agent-id AGENT_ID]
+usage: shaker-agent [-h] [--agent-dir AGENT_DIR] [--agent-id AGENT_ID]
                     [--agent-socket-conn-retries AGENT_SOCKET_CONN_RETRIES]
                     [--agent-socket-recv-timeout AGENT_SOCKET_RECV_TIMEOUT]
                     [--agent-socket-send-timeout AGENT_SOCKET_SEND_TIMEOUT]
@@ -14,6 +14,10 @@ usage: shaker-agent [-h] [--agent-id AGENT_ID]
 
 optional arguments:
   -h, --help            show this help message and exit
+  --agent-dir AGENT_DIR
+                        If specified, directs Shaker to write execution script
+                        for the shell class in agent(s) instance defined
+                        directory. Defaults to /tmp directory.
   --agent-id AGENT_ID   Agent unique id, defaults to MAC of primary interface.
   --agent-socket-conn-retries AGENT_SOCKET_CONN_RETRIES
                         Prior to exiting, the number of reconnects the Agent
diff --git a/doc/source/tools/shaker-all-in-one.txt b/doc/source/tools/shaker-all-in-one.txt
index 155c63e..8122b57 100644
--- a/doc/source/tools/shaker-all-in-one.txt
+++ b/doc/source/tools/shaker-all-in-one.txt
@@ -1,4 +1,5 @@
-usage: shaker-all-in-one [-h] [--agent-join-timeout AGENT_JOIN_TIMEOUT]
+usage: shaker-all-in-one [-h] [--agent-dir AGENT_DIR]
+                         [--agent-join-timeout AGENT_JOIN_TIMEOUT]
                          [--agent-loss-timeout AGENT_LOSS_TIMEOUT]
                          [--artifacts-dir ARTIFACTS_DIR] [--book BOOK]
                          [--cleanup] [--cleanup-on-exit] [--config-dir DIR]
@@ -45,6 +46,10 @@ usage: shaker-all-in-one [-h] [--agent-join-timeout AGENT_JOIN_TIMEOUT]
 
 optional arguments:
   -h, --help            show this help message and exit
+  --agent-dir AGENT_DIR
+                        If specified, directs Shaker to write execution script
+                        for the shell class in agent(s) instance defined
+                        directory. Defaults to /tmp directory.
   --agent-join-timeout AGENT_JOIN_TIMEOUT
                         Timeout to treat agent as join failed in seconds,
                         defaults to env[SHAKER_AGENT_JOIN_TIMEOUT] (time
diff --git a/doc/source/tools/shaker.txt b/doc/source/tools/shaker.txt
index 9600221..725334a 100644
--- a/doc/source/tools/shaker.txt
+++ b/doc/source/tools/shaker.txt
@@ -1,4 +1,5 @@
-usage: shaker [-h] [--agent-join-timeout AGENT_JOIN_TIMEOUT]
+usage: shaker [-h] [--agent-dir AGENT_DIR]
+              [--agent-join-timeout AGENT_JOIN_TIMEOUT]
               [--agent-loss-timeout AGENT_LOSS_TIMEOUT]
               [--artifacts-dir ARTIFACTS_DIR] [--book BOOK]
               [--cleanup-on-exit] [--config-dir DIR] [--config-file PATH]
@@ -32,6 +33,10 @@ usage: shaker [-h] [--agent-join-timeout AGENT_JOIN_TIMEOUT]
 
 optional arguments:
   -h, --help            show this help message and exit
+  --agent-dir AGENT_DIR
+                        If specified, directs Shaker to write execution script
+                        for the shell class in agent(s) instance defined
+                        directory. Defaults to /tmp directory.
   --agent-join-timeout AGENT_JOIN_TIMEOUT
                         Timeout to treat agent as join failed in seconds,
                         defaults to env[SHAKER_AGENT_JOIN_TIMEOUT] (time
diff --git a/etc/shaker.conf b/etc/shaker.conf
index 1368377..e15a285 100644
--- a/etc/shaker.conf
+++ b/etc/shaker.conf
@@ -153,6 +153,11 @@
 # From shaker.engine.config
 #
 
+# If specified, directs Shaker to write execution script for the shell class in
+# agent(s) instance defined directory. Defaults to /tmp directory. (string
+# value)
+#agent_dir = <None>
+
 # Address for server connections (host:port), defaults to
 # env[SHAKER_SERVER_ENDPOINT]. (string value)
 #server_endpoint = <None>
diff --git a/shaker/agent/agent.py b/shaker/agent/agent.py
index 390c851..3b86c74 100644
--- a/shaker/agent/agent.py
+++ b/shaker/agent/agent.py
@@ -61,13 +61,21 @@ def send_reply(socket, agent_id, result):
 def run_command(command):
     command_stdout, command_stderr = None, None
     start = time.time()
+    agent_dir = cfg.CONF.agent_dir
+
+    if agent_dir:
+        utils.mkdir_tree(agent_dir)
 
     if command['type'] == 'program':
         command_stdout, command_stderr = processutils.execute(
             *shlex.split(command['data']), check_exit_code=False)
 
     elif command['type'] == 'script':
-        file_name = tempfile.mktemp()
+        if 'agent_dir' in cfg.CONF:
+            file_name = tempfile.mktemp(dir="%s" % agent_dir)
+        else:
+            file_name = tempfile.mktemp()
+
         with open(file_name, mode='w') as fd:
             fd.write(command['data'])
         LOG.debug('Stored script into %s', file_name)
diff --git a/shaker/engine/config.py b/shaker/engine/config.py
index 93a3988..0c33b2d 100644
--- a/shaker/engine/config.py
+++ b/shaker/engine/config.py
@@ -67,6 +67,10 @@ def generate_output_name():
 
 
 COMMON_OPTS = [
+    cfg.StrOpt('agent-dir', default=utils.env('SHAKER_AGENT_DIR'),
+               help='If specified, directs Shaker to write execution script '
+                    'for the shell class in agent(s) instance defined '
+                    'directory. Defaults to /tmp directory.'),
     cfg.Opt('server-endpoint',
             default=utils.env('SHAKER_SERVER_ENDPOINT'),
             required=True,