eb38b4f373
Reviewers: teon.banek Reviewed By: teon.banek Subscribers: pullbot Differential Revision: https://phabricator.memgraph.io/D2587
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
from argparse import ArgumentParser
|
|
from collections import defaultdict
|
|
import tempfile
|
|
import shutil
|
|
import time
|
|
from common import get_absolute_path, set_cpus
|
|
|
|
try:
|
|
import jail
|
|
except:
|
|
import jail_faker as jail
|
|
|
|
|
|
def wait_for_server(port, delay=0.1):
|
|
cmd = ["nc", "-z", "-w", "1", "127.0.0.1", port]
|
|
while subprocess.call(cmd) != 0:
|
|
time.sleep(0.01)
|
|
time.sleep(delay)
|
|
|
|
|
|
class Memgraph:
|
|
"""
|
|
Knows how to start and stop memgraph.
|
|
"""
|
|
def __init__(self, args, num_workers):
|
|
self.log = logging.getLogger("MemgraphRunner")
|
|
argp = ArgumentParser("MemgraphArgumentParser")
|
|
argp.add_argument("--runner-bin",
|
|
default=get_absolute_path("memgraph", "build"))
|
|
argp.add_argument("--port", default="7687",
|
|
help="Database and client port")
|
|
argp.add_argument("--data-directory", default=None)
|
|
argp.add_argument("--storage-snapshot-on-exit", action="store_true")
|
|
argp.add_argument("--storage-recover-on-startup", action="store_true")
|
|
self.log.info("Initializing Runner with arguments %r", args)
|
|
self.args, _ = argp.parse_known_args(args)
|
|
self.num_workers = num_workers
|
|
self.database_bin = jail.get_process()
|
|
self.name = "memgraph"
|
|
set_cpus("database-cpu-ids", self.database_bin, args)
|
|
|
|
def start(self):
|
|
self.log.info("start")
|
|
database_args = ["--bolt-port", self.args.port,
|
|
"--query-execution-timeout-sec", "0"]
|
|
if self.num_workers:
|
|
database_args += ["--bolt-num-workers", str(self.num_workers)]
|
|
if self.args.data_directory:
|
|
database_args += ["--data-directory", self.args.data_directory]
|
|
if self.args.storage_recover_on_startup:
|
|
database_args += ["--storage-recover-on-startup"]
|
|
if self.args.storage_snapshot_on_exit:
|
|
database_args += ["--storage-snapshot-on-exit"]
|
|
|
|
# find executable path
|
|
runner_bin = self.args.runner_bin
|
|
if not os.path.exists(runner_bin):
|
|
# Apollo builds both debug and release binaries on diff
|
|
# so we need to use the release binary if the debug one
|
|
# doesn't exist
|
|
runner_bin = get_absolute_path("memgraph", "build_release")
|
|
|
|
# start memgraph
|
|
self.database_bin.run(runner_bin, database_args, timeout=600)
|
|
wait_for_server(self.args.port)
|
|
|
|
def stop(self):
|
|
self.database_bin.send_signal(jail.SIGTERM)
|
|
self.database_bin.wait()
|
|
|
|
|
|
class Neo:
|
|
"""
|
|
Knows how to start and stop neo4j.
|
|
"""
|
|
def __init__(self, args, config):
|
|
self.log = logging.getLogger("NeoRunner")
|
|
argp = ArgumentParser("NeoArgumentParser")
|
|
argp.add_argument("--runner-bin", default=get_absolute_path(
|
|
"neo4j/bin/neo4j", "libs"))
|
|
argp.add_argument("--port", default="7687",
|
|
help="Database and client port")
|
|
argp.add_argument("--http-port", default="7474",
|
|
help="Database and client port")
|
|
self.log.info("Initializing Runner with arguments %r", args)
|
|
self.args, _ = argp.parse_known_args(args)
|
|
self.config = config
|
|
self.database_bin = jail.get_process()
|
|
self.name = "neo4j"
|
|
set_cpus("database-cpu-ids", self.database_bin, args)
|
|
|
|
def start(self):
|
|
self.log.info("start")
|
|
|
|
# create home directory
|
|
self.neo4j_home_path = tempfile.mkdtemp(dir="/dev/shm")
|
|
|
|
try:
|
|
os.symlink(os.path.join(get_absolute_path("neo4j", "libs"), "lib"),
|
|
os.path.join(self.neo4j_home_path, "lib"))
|
|
neo4j_conf_dir = os.path.join(self.neo4j_home_path, "conf")
|
|
neo4j_conf_file = os.path.join(neo4j_conf_dir, "neo4j.conf")
|
|
os.mkdir(neo4j_conf_dir)
|
|
shutil.copyfile(self.config, neo4j_conf_file)
|
|
with open(neo4j_conf_file, "a") as f:
|
|
f.write("\ndbms.connector.bolt.listen_address=:" +
|
|
self.args.port + "\n")
|
|
f.write("\ndbms.connector.http.listen_address=:" +
|
|
self.args.http_port + "\n")
|
|
|
|
# environment
|
|
cwd = os.path.dirname(self.args.runner_bin)
|
|
env = {"NEO4J_HOME": self.neo4j_home_path}
|
|
|
|
self.database_bin.run(self.args.runner_bin, args=["console"],
|
|
env=env, timeout=600, cwd=cwd)
|
|
except:
|
|
shutil.rmtree(self.neo4j_home_path)
|
|
raise Exception("Couldn't run Neo4j!")
|
|
|
|
wait_for_server(self.args.http_port, 2.0)
|
|
|
|
def stop(self):
|
|
self.database_bin.send_signal(jail.SIGTERM)
|
|
self.database_bin.wait()
|
|
if os.path.exists(self.neo4j_home_path):
|
|
shutil.rmtree(self.neo4j_home_path)
|