From bf497e91ad1ab31abef16b797c2c756584cbc74e Mon Sep 17 00:00:00 2001
From: anand76 <anand76@devvm4702.ftw0.facebook.com>
Date: Thu, 27 Oct 2022 19:47:01 -0700
Subject: [PATCH] Allow a custom DB cleanup command to be passed to
 db_crashtest.py (#10883)

Summary:
This option allows a custom cleanup command line for a non-Posix file system to be used by db_crashtest.py to cleanup between runs.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/10883

Test Plan: Run the whitebox crash test

Reviewed By: pdillinger

Differential Revision: D40726424

Pulled By: anand1976

fbshipit-source-id: b827f6b583ff78f9ca75ced2d96f7e58f5200432
---
 crash_test.mk         |  2 +-
 tools/db_crashtest.py | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/crash_test.mk b/crash_test.mk
index b7b00a65b..5e8b3573a 100644
--- a/crash_test.mk
+++ b/crash_test.mk
@@ -8,7 +8,7 @@ DB_STRESS_CMD?=./db_stress
 include common.mk
 
 CRASHTEST_MAKE=$(MAKE) -f crash_test.mk
-CRASHTEST_PY=$(PYTHON) -u tools/db_crashtest.py --stress_cmd=$(DB_STRESS_CMD)
+CRASHTEST_PY=$(PYTHON) -u tools/db_crashtest.py --stress_cmd=$(DB_STRESS_CMD) --cleanup_cmd='$(DB_CLEANUP_CMD)'
 
 .PHONY: crash_test crash_test_with_atomic_flush crash_test_with_txn \
 	crash_test_with_best_efforts_recovery crash_test_with_ts \
diff --git a/tools/db_crashtest.py b/tools/db_crashtest.py
index 4c5b10b79..b16643e5a 100644
--- a/tools/db_crashtest.py
+++ b/tools/db_crashtest.py
@@ -210,6 +210,7 @@ _TEST_DIR_ENV_VAR = "TEST_TMPDIR"
 _DEBUG_LEVEL_ENV_VAR = "DEBUG_LEVEL"
 
 stress_cmd = "./db_stress"
+cleanup_cmd = None
 
 
 def is_release_mode():
@@ -224,6 +225,10 @@ def get_dbname(test_name):
     else:
         dbname = test_tmpdir + "/" + test_dir_name
         shutil.rmtree(dbname, True)
+        if cleanup_cmd is not None:
+            print("Running DB cleanup command - %s\n" % cleanup_cmd)
+            # Ignore failure
+            os.system(cleanup_cmd)
         os.mkdir(dbname)
     return dbname
 
@@ -691,6 +696,7 @@ def gen_cmd(params, unknown_params):
                 "write_policy",
                 "stress_cmd",
                 "test_tiered_storage",
+                "cleanup_cmd",
             }
             and v is not None
         ]
@@ -926,6 +932,12 @@ def whitebox_crash_main(args, unknown_args):
             # we need to clean up after ourselves -- only do this on test
             # success
             shutil.rmtree(dbname, True)
+            if cleanup_cmd is not None:
+                print("Running DB cleanup command - %s\n" % cleanup_cmd)
+                ret = os.system(cleanup_cmd)
+                if ret != 0:
+                    print("TEST FAILED. DB cleanup returned error %d\n" % ret)
+                    sys.exit(1)
             os.mkdir(dbname)
             if (expected_values_dir is not None):
                 shutil.rmtree(expected_values_dir, True)
@@ -938,6 +950,7 @@ def whitebox_crash_main(args, unknown_args):
 
 def main():
     global stress_cmd
+    global cleanup_cmd
 
     parser = argparse.ArgumentParser(
         description="This script runs and kills \
@@ -953,6 +966,7 @@ def main():
     parser.add_argument("--write_policy", choices=["write_committed", "write_prepared"])
     parser.add_argument("--stress_cmd")
     parser.add_argument("--test_tiered_storage", action="store_true")
+    parser.add_argument("--cleanup_cmd")
 
     all_params = dict(
         list(default_params.items())
@@ -987,6 +1001,8 @@ def main():
 
     if args.stress_cmd:
         stress_cmd = args.stress_cmd
+    if args.cleanup_cmd:
+        cleanup_cmd = args.cleanup_cmd
     if args.test_type == "blackbox":
         blackbox_crash_main(args, unknown_args)
     if args.test_type == "whitebox":