c-foreach-restart: added option to execute command on shutdown

eevans · eevans · commit 5b03f2f02bf2 · 2016-09-26T19:21:59.000-05:00
diff --git a/c-foreach-restart b/c-foreach-restart
@@ -11,19 +11,22 @@ from cassandra.tools.config import LOG_LEVEL
 logging.basicConfig(level=LOG_LEVEL)
 
 
-def main(retries, delay):
+def main(retries, delay, post_shutdown):
     for instance in get_instances():
-        instance.restart(retries=retries, delay=delay)
+        instance.restart(retries=retries, delay=delay, post_shutdown=post_shutdown)
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Cassandra instance restarter")
     parser.add_argument("-r", "--retries", metavar="RETRIES", default=10, type=int,
                         help="Maximum number of times to check if service is up.")
     parser.add_argument("-d", "--delay", metavar="DELAY", default=6, type=int,
                         help="Number seconds between connection attempts, in seconds.")
+    parser.add_argument("--execute-post-shutdown", metavar="CMD", type=str,
+                        help="Command to execute after Cassandra has been shutdown, and before it "
+                        "is started back up.")
     return parser.parse_args(sys.argv[1:])
 
 
 if __name__ == "__main__":
     args = parse_args()
-    main(args.retries, args.delay)
+    main(args.retries, args.delay, args.execute_post_shutdown)
diff --git a/cassandra/tools/instances.py b/cassandra/tools/instances.py
@@ -44,27 +44,36 @@ def __get(key):
     """
     Restart this Cassandra instance.
     """
-    def restart(self, retries=10, delay=6):
+    def restart(self, retries=10, delay=6, post_shutdown=None):
         self.__log_debug("Restarting instance... (retries=%s, delay=%s)", retries, delay)
         self.__log_info("Disabling client ports...")
         self.nodetool.run("disablebinary")
         self.nodetool.run("disablethrift")
         self.__log_info("Draining...")
         self.nodetool.run("drain")
 
-        # Restart Cassandra
-        self.__log_info("Restarting service %s", self.service_name)
-        (retcode, stdout, stderr) = call("systemctl", "restart", self.service_name)
-        if retcode != 0:
-            self.__log_error("systemctl returned exit code %s", retcode)
+        def exec_command(*args):
+            (retcode, stdout, stderr) = call(*args)
             stdout = stdout.rstrip()
             stderr = stderr.rstrip()
             if stdout:
                 for line in stdout.splitlines():
-                    self.__log_error(line)
+                    self.__log_info(line)
             if stderr:
                 for line in stderr.splitlines():
                     self.__log_error(line)
+            if retcode != 0:
+                self.__log_error("%s returned exit code %s", args[0], retcode)
+                raise RuntimeError("{} returned exit code {}".format(args[0], retcode))
+
+        # Restart Cassandra
+        self.__log_info("Stopping service %s", self.service_name)
+        exec_command("systemctl", "stop", self.service_name)
+        if post_shutdown:
+            self.__log_info("Executing post-shutdown command: %s", post_shutdown)
+            exec_command(*(post_shutdown.strip().split()))
+        self.__log_info("Starting service %s", self.service_name)
+        exec_command("systemctl", "start", self.service_name)
 
         # Wait for Cassandra to come back up before continuing
         listening = False
