openstack-cis/parallel_task_runner.py at master · cloudbase/openstack-cis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import subprocess
from pprint import pprint, pformat
import sys
import time
import os
import signal
from optparse import OptionParser, OptionValueError
import json
import logging
import logging.handlers


loop_check_delay = 5
processes = {}

tasks_example = {
    "Build compute": {
        "cmd": "ansible-playbook -i inventory build-win2016-compute.yaml",
        "log": "hv.log"
    },
    "Build devstack": {
        "cmd": "ansible-playbook -i inventory build-devstack.yaml",
        "log": "dv.log"
    }
}


def roll_log(log_file):
    if os.path.isfile(log_file):
        logger = logging.getLogger('Roll-Logger')
        handler = logging.handlers.RotatingFileHandler(log_file, backupCount=5)
        logger.addHandler(handler)
        logger.handlers[0].doRollover()


# capture SIGINT (ctrl+c) and SIGTERM and kill everything if necessary
def sigint(s, f):
    print("Received signal: %s" % s)
    for t, proc in list(processes.items()):
        print("Killing task: %s" % t)
        os.killpg(os.getpgid(proc["process"].pid), signal.SIGKILL)
    sys.exit(1)


def create_process_element(task_name, task_data, attempt=1):
    roll_log(task_data["log"])
    log = open(task_data["log"], "w")
    result = {
        "process": create_process(task_data["cmd"], log),
        "attempt": attempt,
        "max": int(task_data["retry"]) if "retry" in task_data else 1
    }
    print("Starting task: %s \n    command: %s\n    log: %s\n    attempt: %s of %s" %
          (task_name, task_data["cmd"], task_data["log"], attempt, result["max"]))
    log.close()
    return result


def create_process(cmd, output_log, shell=True, start_new_session=True):
    return subprocess.Popen(
        cmd,
        stdout=output_log,
        stderr=output_log,
        start_new_session=start_new_session,
        shell=shell
    )


def check_tasks_arg(value):
    try:
        t = json.loads(value)
        for t_name, t_data in t.items():
            if all(k in t_data for k in ("log", "cmd")):
                return t
            else:
                print("Some fields are missing in task JSON, working example:")
                pprint(tasks_example)
                sys.exit(1)
    except SystemExit:
        sys.exit(1)
    except:
        raise OptionValueError(
            "invalid value for tasks: \n%s" % (value))


if __name__ == '__main__':
    parser = OptionParser()

    parser.add_option("-t", "--tasks", dest="tasks",
                      help="JSON task list", type=str)

    (options, args) = parser.parse_args()
    tasks = check_tasks_arg(options.tasks)

    signal.signal(signal.SIGINT, sigint)
    signal.signal(signal.SIGTERM, sigint)

    for task_name, task_data in tasks.items():
        processes[task_name] = create_process_element(task_name, task_data)

    print("Checking tasks every %d seconds" % loop_check_delay)

    while processes:
        for task_name, process in list(processes.items()):
            p = process["process"]
            if p.poll() is not None:
                print("Task: %s, attempt %s of %s, finished with code: %s" %
                      (task_name, process["attempt"], process["max"], p.returncode)
                      )

                if p.returncode != 0 and process["attempt"] < process["max"]:
                    processes[task_name] = create_process_element(
                        task_name, tasks[task_name], processes[task_name]["attempt"] + 1
                    )
                else:
                    del processes[task_name]
                    if p.returncode != 0:
                        print("Task %s failed, killing everything" % task_name)
                        for t, proc in list(processes.items()):
                            print("Killing task %s, attempt %s of %s" % (t, proc["attempt"], proc["max"]))
                            os.killpg(os.getpgid(proc["process"].pid), signal.SIGKILL)
                        sys.exit(1)
            time.sleep(loop_check_delay)