235 lines
7.3 KiB
Python
Executable File
235 lines
7.3 KiB
Python
Executable File
#!/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!")
|