veilid/test/scripts/run_local_test.py

174 lines
5.4 KiB
Python
Raw Normal View History

2021-11-22 11:28:30 -05:00
#!/usr/bin/env python3
import sys
import os
import io
import argparse
import subprocess
import signal
import time
from threading import Thread
if sys.version_info < (3, 0, 0):
print(__file__ + ' requires Python 3, while Python ' +
str(sys.version[0] + ' was detected. Terminating. '))
sys.exit(1)
script_dir = os.path.dirname(os.path.realpath(__file__))
veilid_server_exe_debug = os.path.join(script_dir, '..', 'veilid-server',
'target', 'debug', 'veilid-server')
veilid_server_exe_release = os.path.join(
script_dir, '..', 'veilid-server', 'target', 'release', 'veilid-server')
main_process = None
subindex_processes = []
try:
# Python 3, open as binary, then wrap in a TextIOWrapper with write-through.
sys.stdout = io.TextIOWrapper(
open(sys.stdout.fileno(), 'wb', 0), write_through=True)
sys.stderr = io.TextIOWrapper(
open(sys.stderr.fileno(), 'wb', 0), write_through=True)
except TypeError:
# Python 2
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', 0)
def tee(prefix, infile, *files):
"""Print `infile` to `files` in a separate thread."""
def fanout(prefix, infile, *files):
with infile:
for line in iter(infile.readline, b""):
for f in files:
f.write(prefix + line)
f.flush()
t = Thread(target=fanout, args=(prefix, infile,) + files)
t.daemon = True
t.start()
return t
def read_until_local_dial_info(proc, proto):
local_dial_info_str = b"Local Dial Info: "
for ln in iter(proc.stdout.readline, ""):
sys.stdout.buffer.write(ln)
sys.stdout.flush()
idx = ln.find(local_dial_info_str)
if idx != -1:
idx += len(local_dial_info_str)
di = ln[idx:]
if di.startswith(proto):
return di.decode("utf-8").strip()
return None
class CleanChildProcesses:
def __enter__(self):
os.setpgrp() # create new process group, become its leader
def __exit__(self, type, value, traceback):
try:
os.killpg(0, signal.SIGKILL) # kill all processes in my group
except KeyboardInterrupt:
pass
def main():
threads = []
# Parse arguments
parser = argparse.ArgumentParser(description='Run veilid servers locally')
parser.add_argument("count", type=int,
help='number of instances to run')
parser.add_argument("--release", action='store_true',
help='use release mode build')
parser.add_argument("--log_trace", action='store_true',
help='use trace logging')
parser.add_argument("--log_info", action='store_true',
help='use info logging')
parser.add_argument("-w", "--wait-for-debug", action='append',
help='specify subnode index to wait for the debugger')
parser.add_argument("--config-file", type=str,
help='configuration file to specify for the bootstrap node')
parser.add_argument("--protocol", type=bytes, default=b"udp",
help='default protocol to choose for dial info')
args = parser.parse_args()
if args.count < 1:
print("Must specify more than one instance")
sys.exit(1)
veilid_server_exe = None
if args.release:
veilid_server_exe = veilid_server_exe_release
else:
veilid_server_exe = veilid_server_exe_debug
base_args = [veilid_server_exe]
base_args.append("--attach=true")
if args.log_info:
pass
elif args.log_trace:
base_args.append("--trace")
else:
base_args.append("--debug")
if args.config_file:
base_args.append("--config-file={}".format(args.config_file))
# Run primary node and get node id
main_args = base_args.copy()
if args.wait_for_debug and ("0" in args.wait_for_debug):
main_args.append("--wait-for-debug")
print("Running main node: {}".format(str(main_args)))
main_proc = subprocess.Popen(
main_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
print(">>> MAIN NODE PID={}".format(main_proc.pid))
main_di = read_until_local_dial_info(main_proc, args.protocol)
threads.append(
tee(b"Veilid-0: ", main_proc.stdout, open("/tmp/veilid-0-out", "wb"),
getattr(sys.stdout, "buffer", sys.stdout))
)
# Run all secondaries and add primary to bootstrap
for n in range(1, args.count):
# time.sleep(2)
sub_args = base_args.copy()
sub_args.append("--subnode_index={}".format(n))
sub_args.append("--bootstrap={}".format(main_di))
if args.wait_for_debug and (str(n) in args.wait_for_debug):
sub_args.append("--wait-for-debug")
print("Running subnode {}: {}".format(n, str(sub_args)))
sub_proc = subprocess.Popen(
sub_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
print(">>> SUBNODE {} NODE PID={}".format(n, sub_proc.pid))
threads.append(
tee("Veilid-{}: ".format(n).encode("utf-8"), sub_proc.stdout, open("/tmp/veilid-{}-out".format(n), "wb"),
getattr(sys.stdout, "buffer", sys.stdout))
)
for t in threads:
t.join() # wait for IO completion
sys.exit(0)
if __name__ == "__main__":
with CleanChildProcesses():
main()