132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
# Copyright 2021 Memgraph Ltd.
|
|
#
|
|
# Use of this software is governed by the Business Source License
|
|
# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
|
|
# License, and you may not use this file except in compliance with the Business Source License.
|
|
#
|
|
# As of the Change Date specified in that file, in accordance with
|
|
# the Business Source License, use of this software will be governed
|
|
# by the Apache License, Version 2.0, included in the file
|
|
# licenses/APL.txt.
|
|
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from argparse import ArgumentParser
|
|
from collections import defaultdict
|
|
|
|
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
|
|
|
|
# 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)
|