1 files changed, 83 insertions, 0 deletions
diff --git a/src/rocksdb/tools/write_stress_runner.py b/src/rocksdb/tools/write_stress_runner.py
new file mode 100644
index 000000000..f39f79cd4
--- /dev/null
+++ b/src/rocksdb/tools/write_stress_runner.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import argparse
+import random
+
+import subprocess
+import sys
+import time
+
+
+def generate_runtimes(total_runtime):
+    # combination of short runtimes and long runtimes, with heavier
+    # weight on short runtimes
+    possible_runtimes_sec = list(range(1, 10)) + list(range(1, 20)) + [100, 1000]
+    runtimes = []
+    while total_runtime > 0:
+        chosen = random.choice(possible_runtimes_sec)
+        chosen = min(chosen, total_runtime)
+        runtimes.append(chosen)
+        total_runtime -= chosen
+    return runtimes
+
+
+def main(args):
+    runtimes = generate_runtimes(int(args.runtime_sec))
+    print(
+        "Going to execute write stress for " + str(runtimes)
+    )  # noqa: E999 T25377293 Grandfathered in
+    first_time = True
+
+    for runtime in runtimes:
+        kill = random.choice([False, True])
+
+        cmd = "./write_stress --runtime_sec=" + ("-1" if kill else str(runtime))
+
+        if len(args.db) > 0:
+            cmd = cmd + " --db=" + args.db
+
+        if first_time:
+            first_time = False
+        else:
+            # use current db
+            cmd = cmd + " --destroy_db=false"
+        if random.choice([False, True]):
+            cmd = cmd + " --delete_obsolete_files_with_fullscan=true"
+        if random.choice([False, True]):
+            cmd = cmd + " --low_open_files_mode=true"
+
+        print(
+            "Running write_stress for %d seconds (%s): %s"
+            % (runtime, ("kill-mode" if kill else "clean-shutdown-mode"), cmd)
+        )
+
+        child = subprocess.Popen([cmd], shell=True)
+        killtime = time.time() + runtime
+        while not kill or time.time() < killtime:
+            time.sleep(1)
+            if child.poll() is not None:
+                if child.returncode == 0:
+                    break
+                else:
+                    print(
+                        "ERROR: write_stress died with exitcode=%d\n" % child.returncode
+                    )
+                    sys.exit(1)
+        if kill:
+            child.kill()
+        # breathe
+        time.sleep(3)
+
+
+if __name__ == "__main__":
+    random.seed(time.time())
+    parser = argparse.ArgumentParser(
+        description="This script runs and kills \
+        write_stress multiple times"
+    )
+    parser.add_argument("--runtime_sec", default="1000")
+    parser.add_argument("--db", default="")
+    args = parser.parse_args()
+    main(args)