services/keep/tools/traffic_test.py

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