def main():
config = utils.parse_vantage_points()
parser = argparse.ArgumentParser()
parser.add_argument('action', choices=['start', 'stop', 'status'])
parser.add_argument('server_type', choices=['gce_servers', 'emu_servers'])
args = parser.parse_args()
procs = []
for server, server_cfg in config[args.server_type].iteritems():
instance_name = server_cfg['name']
zone = server_cfg['zone']
if args.action == 'start':
procs.append(Popen(
['gcloud', 'compute', 'instances', 'start',
instance_name, '--zone', zone]))
elif args.action == 'stop':
procs.append(Popen(
['gcloud', 'compute', 'instances', 'stop',
instance_name, '--zone', zone]))
elif args.action == 'status':
procs.append(Popen(
['gcloud', 'compute', 'instances', 'describe',
instance_name, '--zone', zone],
stdout=PIPE))
for proc in procs:
if args.action == 'status':
status, _ = proc.communicate()
print(yaml.load(status)['status'])
else:
proc.wait()
def run_without_tunnel(self):
port = utils.get_open_port()
# run the side specified by self.run_first
cmd = ['python', self.cc_src, self.run_first, port]
sys.stderr.write('Running %s %s...\n' % (self.cc, self.run_first))
self.proc_first = Popen(cmd, preexec_fn=os.setsid)
# sleep just in case the process isn't quite listening yet
# the cleaner approach might be to try to verify the socket is open
time.sleep(self.run_first_setup_time)
self.test_start_time = utils.utc_time()
# run the other side specified by self.run_second
sh_cmd = 'python %s %s $MAHIMAHI_BASE %s' % (self.cc_src,
self.run_second, port)
sh_cmd = ' '.join(self.mm_cmd) + " -- sh -c '%s'" % sh_cmd
sys.stderr.write('Running %s %s...\n' % (self.cc, self.run_second))
self.proc_second = Popen(sh_cmd, shell=True, preexec_fn=os.setsid)
signal.signal(signal.SIGALRM, utils.timeout_handler)
signal.alarm(self.runtime)
try:
self.proc_first.wait()
self.proc_second.wait()
except utils.TimeoutError:
pass
else:
signal.alarm(0)
sys.stderr.write('Warning: test exited before time limit\n')
finally:
self.test_end_time = utils.utc_time()
return True
def run_tunnel_managers(self):
# run tunnel server manager
if self.mode == 'remote':
if self.server_side == 'local':
ts_manager_cmd = ['python', self.tunnel_manager]
else:
ts_manager_cmd = self.r['ssh_cmd'] + [
'"python', self.r['tunnel_manager'], '"'
]
else:
ts_manager_cmd = ['python', self.tunnel_manager]
sys.stderr.write('[tunnel server manager (tsm)] ')
self.ts_manager = Popen(ts_manager_cmd,
stdin=PIPE,
stdout=PIPE,
preexec_fn=os.setsid,
shell=True)
ts_manager = self.ts_manager
while True:
running = ts_manager.stdout.readline()
if 'tunnel manager is running' in running:
sys.stderr.write(running)
break
ts_manager.stdin.write('prompt [tsm]\n')
ts_manager.stdin.flush()
# run tunnel client manager
if self.mode == 'remote':
if self.server_side == 'local':
tc_manager_cmd = self.r['ssh_cmd'] + [
'python', self.r['tunnel_manager']
]
else:
tc_manager_cmd = ['python', self.tunnel_manager]
else:
tc_manager_cmd = self.mm_cmd + ['python', self.tunnel_manager]
sys.stderr.write('[tunnel client manager (tcm)] ')
self.tc_manager = Popen(tc_manager_cmd,
stdin=PIPE,
stdout=PIPE,
preexec_fn=os.setsid,
shell=True)
tc_manager = self.tc_manager
while True:
running = tc_manager.stdout.readline()
if 'tunnel manager is running' in running:
sys.stderr.write(running)
break
tc_manager.stdin.write('prompt [tcm]\n')
tc_manager.stdin.flush()
os.system("sudo sysctl -w net.ipv4.ip_forward=1")
return ts_manager, tc_manager
def test_schemes(args):
wrappers_dir = path.join(context.src_dir, 'wrappers')
if args.all:
schemes = utils.parse_config()['schemes'].keys()
elif args.schemes is not None:
schemes = args.schemes.split()
for scheme in schemes:
sys.stderr.write('Testing %s...\n' % scheme)
src = path.join(wrappers_dir, scheme + '.py')
run_first = check_output([src, 'run_first']).strip()
run_second = 'receiver' if run_first == 'sender' else 'sender'
port = utils.get_open_port()
# run first to run
cmd = [src, run_first, port]
first_proc = Popen(cmd, preexec_fn=os.setsid)
# wait for 'run_first' to be ready
time.sleep(3)
# run second to run
cmd = [src, run_second, '127.0.0.1', port]
second_proc = Popen(cmd, preexec_fn=os.setsid)
# test lasts for 3 seconds
signal.signal(signal.SIGALRM, utils.timeout_handler)
signal.alarm(3)
try:
for proc in [first_proc, second_proc]:
proc.wait()
if proc.returncode != 0:
sys.exit('%s failed in tests' % scheme)
except utils.TimeoutError:
pass
except Exception as exception:
sys.exit('test_schemes.py: %s\n' % exception)
else:
signal.alarm(0)
sys.exit('test exited before time limit')
finally:
# cleanup
utils.kill_proc_group(first_proc)
utils.kill_proc_group(second_proc)
def stop_hosts():
if expt_type == 'node':
# run aws.py
aws_py = path.join(context.scripts_dir, 'aws.py')
with open(os.devnull, 'w') as DEVNULL:
Popen([aws_py, 'stop'], stdout=DEVNULL, stderr=DEVNULL).wait()
else:
if expt_type == 'cloud':
gce_server_type = 'gce_servers'
elif expt_type == 'emu':
gce_server_type = 'emu_servers'
# run gce.py
gce_py = path.join(context.scripts_dir, 'gce.py')
with open(os.devnull, 'w') as DEVNULL:
Popen([gce_py, 'stop', gce_server_type],
stdout=DEVNULL,
stderr=DEVNULL).wait()
def execute(host_cmd):
""" Given 'host_cmd' with form: {host: command_to_run_over_SSH, ...},
return {host: boolean (indicating success or failure), ... }
"""
host_proc = {}
for host, cmd in host_cmd.iteritems():
host_proc[host] = Popen(ssh_cmd(host) + [cmd])
for host in host_proc:
host_proc[host] = True if host_proc[host].wait() == 0 else False
return host_proc
def run_tunnel_managers(self):
# run tunnel server manager
if self.mode == 'remote':
if self.server_side == 'local':
ts_manager_cmd = ['python', self.tunnel_manager]
else:
ts_manager_cmd = self.r['ssh_cmd'] + [
'python', self.r['tunnel_manager']
]
else:
ts_manager_cmd = ['python', self.tunnel_manager]
sys.stderr.write('[tunnel server manager (tsm)] ')
self.ts_manager = Popen(ts_manager_cmd,
stdin=PIPE,
stdout=PIPE,
preexec_fn=os.setsid)
ts_manager = self.ts_manager
while True:
running = ts_manager.stdout.readline()
if 'tunnel manager is running' in running:
sys.stderr.write(running)
break
if not running:
sys.stderr.write('WARNING: tunnel server manager terminated '
'unexpectedly\n')
return None, None
ts_manager.stdin.write('prompt [tsm]\n')
ts_manager.stdin.flush()
# run tunnel client manager
if self.mode == 'remote':
if self.server_side == 'local':
tc_manager_cmd = self.r['ssh_cmd'] + [
'python', self.r['tunnel_manager']
]
else:
tc_manager_cmd = ['python', self.tunnel_manager]
else:
tc_manager_cmd = self.mm_cmd + ['python', self.tunnel_manager]
sys.stderr.write('[tunnel client manager (tcm)] ')
# NB: using `preexec_fn=os.setsid` creates a new process group, so that
# it is easy to kill all associated child processes afterwards.
self.tc_manager = Popen(tc_manager_cmd,
stdin=PIPE,
stdout=PIPE,
preexec_fn=os.setsid)
tc_manager = self.tc_manager
while True:
running = tc_manager.stdout.readline()
if 'tunnel manager is running' in running:
sys.stderr.write(running)
break
if not running:
sys.stderr.write('WARNING: tunnel client manager terminated '
'unexpectedly\n')
return ts_manager, None
tc_manager.stdin.write('prompt [tcm]\n')
tc_manager.stdin.flush()
return ts_manager, tc_manager
def start_hosts():
sys.stderr.write('----- Starting hosts -----\n')
# set of live hosts (including nodes and cloud servers)
live_hosts = set()
# list of cloud server to start
clouds_to_start = []
if expt_type == 'node':
host_status = utils.check_ssh_connection(
utils.host_cfg['nodes'].keys())
for host, status in host_status.iteritems():
if status:
live_hosts.add(host)
# get peer cloud of host
node_cfg = utils.host_cfg['nodes'][host]
peer_cloud = node_cfg['peer_cloud']
clouds_to_start.append(peer_cloud)
# run aws.py
aws_py = path.join(context.scripts_dir, 'aws.py')
with open(os.devnull, 'w') as DEVNULL:
Popen([aws_py, 'start', '--hosts', ' '.join(clouds_to_start)],
stdout=DEVNULL,
stderr=DEVNULL).wait()
else:
if expt_type == 'cloud':
gce_server_type = 'gce_servers'
elif expt_type == 'emu':
gce_server_type = 'emu_servers'
for host in utils.host_cfg[gce_server_type]:
clouds_to_start.append(host)
# run gce.py
gce_py = path.join(context.scripts_dir, 'gce.py')
with open(os.devnull, 'w') as DEVNULL:
Popen([gce_py, 'start', gce_server_type],
stdout=DEVNULL,
stderr=DEVNULL).wait()
# check SSH connections to cloud servers
host_status = utils.check_ssh_connection(clouds_to_start,
retry_times=2,
retry_timeout=20)
for host, status in host_status.iteritems():
if not status:
sys.exit('Failed to start %s' % host)
live_hosts.add(host)
# print live hosts
sys.stderr.write('Live hosts:')
if not live_hosts:
sys.stderr.write(' None\n')
else:
for host in live_hosts:
sys.stderr.write(' ' + host)
sys.stderr.write('\n')
return live_hosts
