1#!/usr/bin/env python
2# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
3from __future__ import absolute_import, division, print_function, unicode_literals
4
5import subprocess
6import argparse
7import random
8import time
9import sys
10
11
12def generate_runtimes(total_runtime):
13    # combination of short runtimes and long runtimes, with heavier
14    # weight on short runtimes
15    possible_runtimes_sec = list(range(1, 10)) + list(range(1, 20)) + [100, 1000]
16    runtimes = []
17    while total_runtime > 0:
18        chosen = random.choice(possible_runtimes_sec)
19        chosen = min(chosen, total_runtime)
20        runtimes.append(chosen)
21        total_runtime -= chosen
22    return runtimes
23
24
25def main(args):
26    runtimes = generate_runtimes(int(args.runtime_sec))
27    print("Going to execute write stress for " + str(runtimes))  # noqa: E999 T25377293 Grandfathered in
28    first_time = True
29
30    for runtime in runtimes:
31        kill = random.choice([False, True])
32
33        cmd = './write_stress --runtime_sec=' + \
34            ("-1" if kill else str(runtime))
35
36        if len(args.db) > 0:
37            cmd = cmd + ' --db=' + args.db
38
39        if first_time:
40            first_time = False
41        else:
42            # use current db
43            cmd = cmd + ' --destroy_db=false'
44        if random.choice([False, True]):
45            cmd = cmd + ' --delete_obsolete_files_with_fullscan=true'
46        if random.choice([False, True]):
47            cmd = cmd + ' --low_open_files_mode=true'
48
49        print("Running write_stress for %d seconds (%s): %s" %
50              (runtime, ("kill-mode" if kill else "clean-shutdown-mode"),
51              cmd))
52
53        child = subprocess.Popen([cmd], shell=True)
54        killtime = time.time() + runtime
55        while not kill or time.time() < killtime:
56            time.sleep(1)
57            if child.poll() is not None:
58                if child.returncode == 0:
59                    break
60                else:
61                    print("ERROR: write_stress died with exitcode=%d\n"
62                          % child.returncode)
63                    sys.exit(1)
64        if kill:
65            child.kill()
66        # breathe
67        time.sleep(3)
68
69if __name__ == '__main__':
70    random.seed(time.time())
71    parser = argparse.ArgumentParser(description="This script runs and kills \
72        write_stress multiple times")
73    parser.add_argument("--runtime_sec", default='1000')
74    parser.add_argument("--db", default='')
75    args = parser.parse_args()
76    main(args)
77