#!/usr/bin/env python3 # -*- coding: utf-8 -*- import argparse import copy import multiprocessing import os import subprocess import time from argparse import Namespace as Args from subprocess import Popen from typing import Dict, List from test_config import LARGE_DATASET, SMALL_DATASET, DatabaseMode, DatasetConstants # paths SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__)) BASE_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..")) BUILD_DIR = os.path.join(BASE_DIR, "build") MEASUREMENTS_FILE = os.path.join(SCRIPT_DIR, ".apollo_measurements") KEY_FILE = os.path.join(SCRIPT_DIR, ".key.pem") CERT_FILE = os.path.join(SCRIPT_DIR, ".cert.pem") # long running stats file STATS_FILE = os.path.join(SCRIPT_DIR, ".long_running_stats") # get number of threads THREADS = os.environ["THREADS"] if "THREADS" in os.environ else multiprocessing.cpu_count() def generate_temporary_ssl_certs(): # https://unix.stackexchange.com/questions/104171/create-ssl-certificate-non-interactively subj = "/C=HR/ST=Zagreb/L=Zagreb/O=Memgraph/CN=db.memgraph.com" subprocess.run( [ "openssl", "req", "-new", "-newkey", "rsa:4096", "-days", "365", "-nodes", "-x509", "-subj", subj, "-keyout", KEY_FILE, "-out", CERT_FILE, ], check=True, ) def remove_certificates() -> None: os.remove(KEY_FILE) os.remove(CERT_FILE) def parse_arguments() -> Args: # parse arguments parser = argparse.ArgumentParser(description="Run stress tests on Memgraph.") parser.add_argument("--memgraph", default=os.path.join(BUILD_DIR, "memgraph")) parser.add_argument("--log-file", default="") parser.add_argument("--data-directory", default="") parser.add_argument("--python", default=os.path.join(SCRIPT_DIR, "ve3", "bin", "python3"), type=str) parser.add_argument("--large-dataset", action="store_const", const=True, default=False) parser.add_argument("--small-dataset", action="store_const", const=True, default=False) parser.add_argument("--specific-test", default="") parser.add_argument("--use-ssl", action="store_const", const=True, default=False) parser.add_argument("--verbose", action="store_const", const=True, default=False) return parser.parse_args() def wait_for_server(port, delay=0.1) -> None: cmd = ["nc", "-z", "-w", "1", "127.0.0.1", str(port)] while subprocess.call(cmd) != 0: time.sleep(0.01) time.sleep(delay) def start_memgraph(args: Args) -> Popen: """Starts Memgraph and return the process""" cwd = os.path.dirname(args.memgraph) cmd = [ args.memgraph, "--bolt-num-workers=" + str(THREADS), "--storage-properties-on-edges=true", "--storage-snapshot-on-exit=true", "--storage-snapshot-interval-sec=600", "--storage-snapshot-retention-count=1", "--storage-wal-enabled=true", "--storage-recover-on-startup=false", "--query-execution-timeout-sec=1200", "--bolt-server-name-for-init=Neo4j/", ] if not args.verbose: cmd += ["--log-level", "WARNING"] if args.log_file: cmd += ["--log-file", args.log_file] if args.data_directory: cmd += ["--data-directory", args.data_directory] if args.use_ssl: cmd += ["--bolt-cert-file", CERT_FILE, "--bolt-key-file", KEY_FILE] memgraph_proc = subprocess.Popen(cmd, cwd=cwd) wait_for_server(7687) assert memgraph_proc.poll() is None, "The database binary died prematurely!" return memgraph_proc def stop_memgraph(proc_mg: Popen) -> None: proc_mg.terminate() ret_mg = proc_mg.wait() if ret_mg != 0: raise Exception("Memgraph binary returned non-zero ({})!".format(ret_mg)) def run_test(args: Args, test: str, options: List[str], timeout: int, mode: DatabaseMode) -> float: """Runs tests for a set of specific database configuration. Args: args: Arguments passed to the test test: Test name options: List of options specific for each test timeout: Timeout in minutes mode: DatabaseMode (storage mode & isolation level pair) """ print("Running test '{}'".format(test)) binary = _find_test_binary(args, test) # start test cmd = ( binary + [ "--worker-count", str(THREADS), "--isolation-level", mode.isolation_level, "--storage-mode", mode.storage_mode, ] + options ) start = time.time() ret_test = subprocess.run(cmd, cwd=SCRIPT_DIR, timeout=timeout * 60) if ret_test.returncode != 0: raise Exception("Test '{}' binary returned non-zero ({})!".format(test, ret_test.returncode)) runtime = time.time() - start print(" Done after {:.3f} seconds".format(runtime)) return runtime def _find_test_binary(args: Args, test: str) -> List[str]: if test.endswith(".py"): logging = "DEBUG" if args.verbose else "WARNING" return [args.python, "-u", os.path.join(SCRIPT_DIR, test), "--logging", logging] if test.endswith(".cpp"): exe = os.path.join(BUILD_DIR, "tests", "stress", test[:-4]) return [exe] raise Exception("Test '{}' binary not supported!".format(test)) def run_stress_test_suite(args: Args) -> Dict[str, float]: def cleanup(memgraph_proc): if memgraph_proc.poll() != None: return memgraph_proc.kill() memgraph_proc.wait() runtimes = {} if args.large_dataset and args.small_dataset: raise Exception("Choose only a large dataset or a small dataset for stress testing!") dataset = LARGE_DATASET if args.large_dataset else SMALL_DATASET if args.specific_test: dataset = [x for x in dataset if x[DatasetConstants.TEST] == args.specific_test] if not len(dataset): raise Exception("Specific dataset is not found for stress testing!") for test in dataset: if args.use_ssl: test[DatasetConstants.OPTIONS] += ["--use-ssl"] for mode in test[DatasetConstants.MODE]: test_run = copy.deepcopy(test) # Run for every specified combination of storage mode, serialization mode, etc (if extended) test_run[DatasetConstants.MODE] = mode memgraph_proc = start_memgraph(args) runtime = run_test(args, **test_run) runtimes[os.path.splitext(test[DatasetConstants.TEST])[0]] = runtime stop_memgraph(memgraph_proc) cleanup(memgraph_proc) return runtimes def write_stats(runtimes: Dict[str, float]) -> None: measurements = "" for key, value in runtimes.items(): measurements += "{}.runtime {}\n".format(key, value) with open(STATS_FILE) as f: stats = f.read().split("\n") measurements += "long_running.queries.executed {}\n".format(stats[0]) measurements += "long_running.queries.failed {}\n".format(stats[1]) with open(MEASUREMENTS_FILE, "w") as f: f.write(measurements) if __name__ == "__main__": args = parse_arguments() if args.use_ssl: generate_temporary_ssl_certs() runtimes = run_stress_test_suite(args) if args.use_ssl: remove_certificates() write_stats(runtimes) print("Successfully ran stress tests!")