services/keep/tools/traffic_test.py

   1 #! /usr/bin/env python
   2
   3 # traffic_test.py
   4 #
   5 # Launch a test Keep and API server and PUT and GET a bunch of blocks.
   6 # Can be used to simulate client traffic in Keep to evaluate memory usage,
   7 # error logging, performance, etc.
   8 #
   9 # This script is warty and is relatively environment-specific, but the
  10 # example run described below should execute cleanly.
  11 #
  12 # Usage:
  13 #   traffic_test.py start
  14 #       Starts the test servers.
  15 #   traffic_test.py put file1 file2 file3 ....
  16 #       Runs arv-put on each file.
  17 #   traffic_test.py get hash1 hash2 hash3 ....
  18 #       Loops forever issuing GET requests for specified blocks.
  19 #   traffic_test.py stop
  20 #       Stops the test servers.
  21 #
  22 # Example:
  23 #
  24 #   $ ./traffic_test.py start
  25 #   $ ./traffic_test.py put GS00253-DNA_A02_200_37.tsv.bz2 \
  26 #         GS00253-DNA_B01_200_37.tsv.bz2 \
  27 #         GS00253-DNA_B02_200_37.tsv.bz2
  28 #   $ ./traffic_test.py get $(find /tmp/tmp* -type f -printf "%f ")
  29 #     [loops forever]
  30 #     ^C
  31 #   $ ./traffic_test.py stop
  32 #
  33 # Multiple "get" runs may be run concurrently to evaluate Keep's handling
  34 # of additional concurrent clients.
  35
  36 PYSDK_DIR    = "../../../sdk/python"
  37 PYTEST_DIR   = PYSDK_DIR + "/tests"
  38 ARV_PUT_PATH = PYSDK_DIR + "/bin/arv-put"
  39 ARV_GET_PATH = PYSDK_DIR + "/bin/arv-get"
  40 SECONDS_BETWEEN_GETS = 1
  41
  42 import argparse
  43 import httplib2
  44 import os
  45 import random
  46 import subprocess
  47 import sys
  48 import time
  49
  50 # for run_test_server.py
  51 sys.path.insert(0, PYSDK_DIR)
  52 sys.path.insert(0, PYTEST_DIR)
  53 import arvados
  54 import run_test_server
  55
  56 def arv_cmd(*args):
  57     p = subprocess.Popen([sys.executable] + list(args),
  58                          stdout=subprocess.PIPE)
  59     (arvout, arverr) = p.communicate()
  60     if p.returncode != 0:
  61         print "error {} from {} {}: {}".format(
  62             p.returncode, sys.executable, args, arverr)
  63         sys.exit(p.returncode)
  64     return arvout
  65
  66 def start():
  67     run_test_server.run()
  68     run_test_server.run_keep()
  69
  70 def put(files):
  71     os.environ["ARVADOS_API_HOST"] = "127.0.0.1:3001"
  72     run_test_server.authorize_with('active')
  73     for v in ["ARVADOS_API_HOST",
  74               "ARVADOS_API_HOST_INSECURE",
  75               "ARVADOS_API_TOKEN"]:
  76         os.environ[v] = arvados.config.settings()[v]
  77
  78     if not os.environ.has_key('PYTHONPATH'):
  79         os.environ['PYTHONPATH'] = ''
  80     os.environ['PYTHONPATH'] = "{}:{}:{}".format(
  81         PYSDK_DIR, PYTEST_DIR, os.environ['PYTHONPATH'])
  82
  83     for c in files:
  84         manifest_uuid = arv_cmd(ARV_PUT_PATH, c)
  85
  86 def get(blocks):
  87     os.environ["ARVADOS_API_HOST"] = "127.0.0.1:3001"
  88
  89     run_test_server.authorize_with('active')
  90     for v in ["ARVADOS_API_HOST",
  91               "ARVADOS_API_HOST_INSECURE",
  92               "ARVADOS_API_TOKEN"]:
  93         os.environ[v] = arvados.config.settings()[v]
  94
  95     nqueries = 0
  96     while True:
  97         b = random.choice(blocks)
  98         print "GET /" + b
  99         body = arv_cmd(ARV_GET_PATH, b)
 100         print "got {} bytes".format(len(body))
 101         time.sleep(SECONDS_BETWEEN_GETS)
 102         nqueries = nqueries + 1
 103
 104 def stop():
 105     run_test_server.stop_keep()
 106     run_test_server.stop()
 107
 108 if __name__ == "__main__":
 109     parser = argparse.ArgumentParser()
 110     parser.add_argument('action',
 111                         type=str,
 112                         nargs='+',
 113                         help='''"start", "put", "get", "stop"''')
 114     args = parser.parse_args()
 115
 116     if args.action[0] == 'start':
 117         start()
 118     elif args.action[0] == 'put':
 119         put(args.action[1:])
 120     elif args.action[0] == 'get':
 121         get(args.action[1:])
 122     elif args.action[0] == 'stop':
 123         stop()
 124     else:
 125         print('Unrecognized action "{}"'.format(args.action))
 126         print('actions are "start", "put", "get", "stop"')