From 18d50be4a806c9e0c2c95d3a304a09782858bddc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Sun, 9 Dec 2018 14:12:56 +0100
Subject: [PATCH] add new script to run commands on servers

This script provides parallel remote execution of commands, while having
some special knownledge of servers that should *not* be handled in parallel.

It defers terminal-handling to tmux(1).

It has some targeting capacities using keywords. Commas for 'OR' and slashes
for 'AND', ex: ext/test,saas/test/passerelle will select all external test
servers + all passerelle servers on the SaaS.

It takes any shell command and has some builtin shortcuts such as apt.update
and apt.upgrade. (that's the whole lot, actually).

Regarding actual performance benefits, apt upgrade with no packages to
upgrade:

 $ time eoptasks -k ext/test apt.upgrade
real    0m24,249s
user    0m0,140s
sys     0m0,025s

 $ time eotasks -g ext_test apt.upgrade
real    6m9,956s
user    3m32,096s
sys     0m2,322s
---
 eoptasks/eoptasks.py | 106 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 106 insertions(+)
 create mode 100755 eoptasks/eoptasks.py

diff --git a/eoptasks/eoptasks.py b/eoptasks/eoptasks.py
new file mode 100755
index 0000000..4df8bb2
--- /dev/null
+++ b/eoptasks/eoptasks.py
@@ -0,0 +1,106 @@
+#! /usr/bin/python3
+#
+# This script provides parallel remote execution of commands, while having
+# some special knownledge of servers that should *not* be handled in parallel.
+#
+# It defers terminal-handling to tmux(1).
+#
+# It has some targeting capacities using keywords. Commas for 'OR' and slashes
+# for 'AND', ex: ext/test,saas/test/passerelle will select all external test
+# servers + all passerelle servers on the SaaS.
+#
+# It takes any shell command and has some builtin shortcuts such as apt.update
+# and apt.upgrade. (that's the whole lot, actually).
+#
+# Requirements: libtmux and pyyaml.
+
+import argparse
+import os
+import random
+import re
+import sys
+import time
+
+import libtmux
+import yaml
+
+class Server:
+    def __init__(self, servername, group=''):
+        self.name = servername
+        self.keywords = set(re.split(r'[-_ \.]', servername + ' ' + group))
+
+    def __repr__(self):
+        return '<Server %s %r>' % (self.name, self.keywords)
+
+    def shell(self):
+        return 'ssh %s' % self.name
+
+    def cmd(self, cmd):
+        return 'ssh %s "%s"' % (self.name, cmd)
+
+
+servers = []
+
+servergroups = yaml.load(open('/home/fred/src/eo/puppet/data/servergroups.yaml'))['servergroups']
+for group in servergroups:
+    for servername in servergroups[group]:
+        servers.append(Server(servername, group))
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-k', dest='keywords', type=str)
+parser.add_argument('cmd', type=str)
+
+args = parser.parse_args()
+
+selected_servers = []
+if args.keywords:
+    for keyword in args.keywords.split(','):
+        keywords = set(keyword.split('/'))
+        selected_servers.extend([
+            x for x in servers
+            if keywords.issubset(x.keywords) and not x in selected_servers])
+else:
+    selected_servers = servers
+
+if not selected_servers:
+    sys.exit(0)
+
+cmd = {
+    'apt.update': 'sudo apt update',
+    'apt.upgrade': 'sudo apt update && sudo apt upgrade -y',
+}.get(args.cmd, args.cmd)
+
+tmux_session_name = 's%s' % random.randrange(1000)
+os.system('tmux new-session -s %s -d /bin/bash -c "sleep 2"' % tmux_session_name)
+tmux = libtmux.Server()
+session = tmux.find_where({'session_name': tmux_session_name})
+
+pid = os.fork()
+if pid:
+    os.system('tmux attach-session -t %s' % tmux_session_name)
+else:
+    def cluster_name(server_name):
+        return re.match(r'(.*?)(\d*)$', server_name).group(1).replace(
+                '.rbx.', '.loc.').replace('.gra.', '.loc.').replace('.sbg.', '.loc.')
+
+    while selected_servers:
+        current_clusters = [cluster_name(x.name) for x in session.list_windows()]
+        for server in selected_servers[:]:
+            if cluster_name(server.name) in current_clusters:
+                continue
+            selected_servers.remove(server)
+            session.new_window(
+                    attach=False,
+                    window_name=server.name,
+                    window_shell=server.cmd(cmd))
+            break
+        else:
+            time.sleep(0.1)
+            if len(session.list_windows()) == 1:
+                # a single window remains but there are still unparallelizable
+                # servers to process, create temporary tmux window to avoid tmux
+                # being destroyed as the last window process is over.
+                session.new_window(attach=False, window_name='sleep',
+                        window_shell='/bin/bash -c "sleep 2"')
+        while len(session.list_windows()) > 5:
+            time.sleep(0.1)